# coding=utf-8 # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Runtime monkey-patches making vLLM v1 scheduler CFG-pair-aware. Call ``apply_cfg_patches()`` **before** creating the ``LLM`` instance. Post-sampling token synchronization is handled separately by ``CFGLogitsProcessor._ensure_sample_patched()`` (runs in each worker). """ from __future__ import annotations import logging from typing import Optional logger = logging.getLogger(__name__) _applied = False def _get_cfg_pair_id(request) -> Optional[str]: sp = request.sampling_params if sp and sp.extra_args: return sp.extra_args.get("cfg_pair_id") return None def _get_cfg_role(request) -> Optional[str]: sp = request.sampling_params if sp and sp.extra_args: return sp.extra_args.get("cfg_role") return None def _hold_incomplete_pairs(scheduler) -> list: """Remove from the waiting queue any CFG request whose pair partner has not yet been admitted to the engine. vLLM v1's EngineCore runs a schedule step as soon as the first request of a pair arrives over IPC; for small/fast models the partner often arrives only after that step, so the lone member prefills one step early and the pair is permanently offset by one token (breaking CFG). Holding the lone member until both are present guarantees they prefill in the same step. """ held: list = [] keep: list = [] for req in list(scheduler.waiting): pair_id = scheduler._cfg_req_to_pair.get(req.request_id) complete = True if pair_id is not None: roles = scheduler._cfg_pairs.get(pair_id, {}) complete = len(roles) == 2 and all( rid in scheduler.requests for rid in roles.values() ) (keep if complete else held).append(req) if held: scheduler.waiting.clear() scheduler.waiting.extend(keep) return held def _reorder_waiting_for_cfg(scheduler) -> None: """Move CFG pair partners adjacent in the waiting queue.""" waiting = scheduler.waiting if len(waiting) < 2: return requests = list(waiting) waiting.clear() seen: set[str] = set() result: list = [] for req in requests: rid = req.request_id if rid in seen: continue seen.add(rid) result.append(req) pair_id = scheduler._cfg_req_to_pair.get(rid) if pair_id is None: continue roles = scheduler._cfg_pairs.get(pair_id, {}) for _role, partner_id in roles.items(): if partner_id != rid and partner_id not in seen: for r in requests: if r.request_id == partner_id: seen.add(partner_id) result.append(r) break waiting.extend(result) def _equalize_cfg_pair_progress(scheduler, scheduler_output) -> None: """Ensure both members of every CFG pair reach the same num_computed_tokens. After ``_orig_schedule`` (which calls ``_update_after_schedule``), ``num_computed_tokens`` is already advanced. Prefix-cache hits or unequal chunked-prefill budget can leave one member ahead of the other, causing it to finish prefill a step earlier and emit a decode token the partner misses. Fix: reduce the faster member's allocation so both land on the same ``num_computed_tokens``. The over-allocated KV-cache blocks remain and are consumed in the next step — no blocks are leaked. """ for _pair_id, roles in scheduler._cfg_pairs.items(): cond_id = roles.get("cond") uncond_id = roles.get("uncond") if not cond_id or not uncond_id: continue cond_sched = scheduler_output.num_scheduled_tokens.get(cond_id, 0) uncond_sched = scheduler_output.num_scheduled_tokens.get(uncond_id, 0) if cond_sched == 0 or uncond_sched == 0: continue cond_req = scheduler.requests.get(cond_id) uncond_req = scheduler.requests.get(uncond_id) if not cond_req or not uncond_req: continue if cond_req.num_computed_tokens == uncond_req.num_computed_tokens: continue target = min(cond_req.num_computed_tokens, uncond_req.num_computed_tokens) feasible = True for req, sched in [(cond_req, cond_sched), (uncond_req, uncond_sched)]: if req.num_computed_tokens - target >= sched: feasible = False break if not feasible: continue for req_id, req, orig_sched in [ (cond_id, cond_req, cond_sched), (uncond_id, uncond_req, uncond_sched), ]: diff = req.num_computed_tokens - target if diff > 0: req.num_computed_tokens = target scheduler_output.num_scheduled_tokens[req_id] = orig_sched - diff scheduler_output.total_num_scheduled_tokens -= diff logger.debug( "CFG equalize %s: %d -> %d scheduled, computed -> %d", req_id, orig_sched, orig_sched - diff, target, ) def _co_preempt_split_pairs(scheduler, scheduler_output) -> None: """Detect a split CFG pair (one running, one not / desynced). On vLLM 0.20 the original manual re-preemption surgery is unsafe, so this is a detection-and-log guard; ``_hold_incomplete_pairs`` prevents the split. """ if not scheduler._cfg_pairs: return running_ids = {r.request_id for r in scheduler.running} to_preempt: list[str] = [] for _pair_id, roles in scheduler._cfg_pairs.items(): cond_id = roles.get("cond") uncond_id = roles.get("uncond") if cond_id is None or uncond_id is None: continue if cond_id not in scheduler.requests or uncond_id not in scheduler.requests: continue cond_running = cond_id in running_ids uncond_running = uncond_id in running_ids if cond_running != uncond_running: to_preempt.append(cond_id if cond_running else uncond_id) continue if cond_running and uncond_running: cond_req = scheduler.requests.get(cond_id) uncond_req = scheduler.requests.get(uncond_id) if (cond_req and uncond_req and cond_req.num_computed_tokens != uncond_req.num_computed_tokens): faster = cond_id if cond_req.num_computed_tokens > uncond_req.num_computed_tokens else uncond_id to_preempt.append(faster) if to_preempt: logger.warning( "CFG pair split detected (to_preempt=%s) but co-preempt surgery is " "disabled on vLLM 0.20; relying on pair-hold + equalize.", to_preempt, ) return def _patch_scheduler() -> None: from vllm.v1.core.sched.scheduler import Scheduler _orig_init = Scheduler.__init__ _orig_add = Scheduler.add_request _orig_finish = Scheduler.finish_requests _orig_schedule = Scheduler.schedule def _init(self, *args, **kwargs): _orig_init(self, *args, **kwargs) self._cfg_pairs: dict[str, dict[str, str]] = {} self._cfg_req_to_pair: dict[str, str] = {} def _add(self, request): _orig_add(self, request) pair_id = _get_cfg_pair_id(request) role = _get_cfg_role(request) if pair_id and role: self._cfg_pairs.setdefault(pair_id, {})[role] = request.request_id self._cfg_req_to_pair[request.request_id] = pair_id def _finish(self, request_ids, finished_status): if request_ids is None: result = _orig_finish(self, request_ids, finished_status) self._cfg_pairs.clear() self._cfg_req_to_pair.clear() return result if isinstance(request_ids, str): request_ids = {request_ids} else: request_ids = set(request_ids) partner_ids: set[str] = set() for req_id in request_ids: pair_id = self._cfg_req_to_pair.get(req_id) if pair_id is None: continue for _role, rid in self._cfg_pairs.get(pair_id, {}).items(): if rid != req_id and rid in self.requests: partner_ids.add(rid) all_ids = request_ids | partner_ids result = _orig_finish(self, all_ids, finished_status) for req_id in all_ids: pair_id = self._cfg_req_to_pair.pop(req_id, None) if pair_id: self._cfg_pairs.pop(pair_id, None) return result def _schedule(self): if not self._cfg_pairs: return _orig_schedule(self) _reorder_waiting_for_cfg(self) held = _hold_incomplete_pairs(self) orig_threshold = self.scheduler_config.long_prefill_token_threshold self.scheduler_config.long_prefill_token_threshold = ( self.max_num_scheduled_tokens // 2 ) scheduler_output = _orig_schedule(self) self.scheduler_config.long_prefill_token_threshold = orig_threshold for req in reversed(held): self.waiting.prepend_request(req) _equalize_cfg_pair_progress(self, scheduler_output) _co_preempt_split_pairs(self, scheduler_output) return scheduler_output Scheduler.__init__ = _init Scheduler.add_request = _add Scheduler.finish_requests = _finish Scheduler.schedule = _schedule def apply_cfg_patches() -> None: """Apply CFG scheduler monkey-patches to vLLM v1. Call before ``LLM()``. Token synchronization (post-sampling copy of cond token to uncond slot) is handled by ``CFGLogitsProcessor._ensure_sample_patched()`` which runs inside each worker process automatically. """ global _applied if _applied: return _applied = True logger.info("Applying CFG scheduler patches to vLLM v1") _patch_scheduler()