# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # # Adapted from nvidia/Nemotron-Labs-Audex-2B (Apache-2.0): # inference_scripts_vllm/audiogen_scripts/vllm_cfg_patch.py # Divergences from the official file are commented at the divergence site. """Model-neutral classifier-free-guidance request pairing for the vLLM v1 scheduler. A guided request decodes together with an unconditional companion so their logits can be blended every step. Blending is only correct when both members decode the same position in the same engine step, so :func:`apply_cfg_patches` makes the scheduler pair-aware: a lone member waits for its partner, partners stay adjacent in the waiting queue, their ``num_computed_tokens`` are equalized after every schedule, and they finish together. Requests opt in through ``SamplingParams.extra_args``:: cond: {"cfg_role": "cond", "cfg_pair_id": } uncond: {"cfg_role": "uncond", "cfg_pair_id": } ``cfg_pair_id`` is optional. A model whose serving layer cannot mint a pair id (the companion request id is only known inside the engine) may send ``cfg_role`` alone; the pair id is then derived from the request id, stripping the companion's registered suffix (see :func:`register_cfg_uncond_suffix`). Logit blending itself is the caller's concern: this module only keeps the two rows step-locked. The patch is a no-op for engines that never see a CFG request (the pair registry stays empty). """ from __future__ import annotations from collections import deque from typing import Any from vllm.logger import init_logger logger = init_logger(__name__) __all__ = [ "apply_cfg_patches", "cfg_patches_applied", "cfg_uncond_suffixes", "register_cfg_uncond_suffix", ] _COND = "cond" _UNCOND = "uncond" # Request-id suffix appended to the companion request by the Audex stage input # processor (``CFG_UNCOND_SUFFIX`` in # vllm_omni/model_executor/stage_input_processors/audex.py). Models that mint # the companion request id differently register their own suffix. _UNCOND_REQUEST_ID_SUFFIXES: set[str] = {"__cfg_uncond"} def register_cfg_uncond_suffix(suffix: str) -> None: """Register a companion request-id suffix used to derive a missing pair id.""" if not suffix or not isinstance(suffix, str): raise ValueError(f"CFG uncond request-id suffix must be a non-empty string, got {suffix!r}") _UNCOND_REQUEST_ID_SUFFIXES.add(suffix) def cfg_uncond_suffixes() -> frozenset[str]: """Companion request-id suffixes registered in this process.""" return frozenset(_UNCOND_REQUEST_ID_SUFFIXES) def _derive_pair_id(request_id: str, role: str) -> str: """Pair id for a request that carries ``cfg_role`` but no ``cfg_pair_id``. The cond request id IS the pair id; its companion carries the same id plus a registered suffix, so stripping the suffix recovers it. Longest suffix first, so a registered suffix extending another still resolves. """ if role == _UNCOND: for suffix in sorted(_UNCOND_REQUEST_ID_SUFFIXES, key=len, reverse=True): if request_id.endswith(suffix): return request_id[: -len(suffix)] return request_id def _cfg_extra(request: Any, key: str) -> Any: sampling_params = getattr(request, "sampling_params", None) extra_args = getattr(sampling_params, "extra_args", None) if extra_args: return extra_args.get(key) return None def _wait_queues(scheduler: Any) -> list[Any]: # Divergence from the official (vLLM 0.20-era) file: vLLM 0.24 parks # structurally not-ready requests in a second ``skipped_waiting`` queue; # pair-hold must cover both queues or a lone member parked there would # be admitted without its partner. queues = [scheduler.waiting] skipped = getattr(scheduler, "skipped_waiting", None) if skipped is not None: queues.append(skipped) return queues def _pair_complete(scheduler: Any, request_id: str, blocked_ids: set[str] | None = None) -> bool: pair_id = scheduler._cfg_req_to_pair.get(request_id) if pair_id is None: return True roles = scheduler._cfg_pairs.get(pair_id, {}) if len(roles) != 2 or not all(rid in scheduler.requests for rid in roles.values()): return False if blocked_ids: # A partner parked in skipped_waiting (e.g. waiting for remote KVs) # may fail promotion this step while this member gets scheduled from # waiting; hold this member until the partner is schedulable. for rid in roles.values(): if rid != request_id and rid in blocked_ids: return False return True # Schedule steps a lone pair member may wait for its partner. Partners arrive # within a handful of steps in practice; exhausting this budget means the # companion was never created (e.g. prompt expansion failed), and the request # is released to decode unguided instead of hanging forever. _MAX_PAIR_HOLD_STEPS = 512 def _drop_stale_pair(scheduler: Any, request_id: str) -> None: pair_id = scheduler._cfg_req_to_pair.get(request_id) if pair_id is None: return for rid in scheduler._cfg_pairs.pop(pair_id, {}).values(): scheduler._cfg_req_to_pair.pop(rid, None) logger.warning( "CFG pair %s: partner of %s never arrived after %d schedule steps; releasing the request unguided", pair_id, request_id, _MAX_PAIR_HOLD_STEPS, ) def _hold_incomplete_pairs(scheduler: Any) -> list[tuple[Any, Any]]: """Pull CFG requests whose partner is not yet admitted out of the wait queues. The engine schedules as soon as the first member of a pair arrives over IPC; without the hold, the lone member prefills one step early and the pair is permanently offset by one token. Held requests are re-prepended after the schedule step. """ hold_counts: dict[str, int] = getattr(scheduler, "_cfg_hold_counts", None) or {} scheduler._cfg_hold_counts = hold_counts skipped = getattr(scheduler, "skipped_waiting", None) blocked_ids = {req.request_id for req in list(skipped)} if skipped is not None else set() held: list[tuple[Any, Any]] = [] for queue in _wait_queues(scheduler): # A member is only "blocked" for its partner's sake when it sits in # the OTHER queue; members of the queue being scanned move together. partner_blockers = blocked_ids if queue is not skipped else set() held_here = [] for request in list(queue): if _pair_complete(scheduler, request.request_id, partner_blockers): hold_counts.pop(request.request_id, None) continue count = hold_counts.get(request.request_id, 0) + 1 if count > _MAX_PAIR_HOLD_STEPS: hold_counts.pop(request.request_id, None) _drop_stale_pair(scheduler, request.request_id) continue hold_counts[request.request_id] = count held_here.append(request) if held_here: queue.remove_requests(held_here) held.extend((queue, req) for req in held_here) return held def _release_held(held: list[tuple[Any, Any]]) -> None: for queue, request in reversed(held): queue.prepend_request(request) def _reorder_waiting_for_cfg(scheduler: Any) -> None: """Move CFG pair partners adjacent in the FCFS waiting queue.""" waiting = scheduler.waiting # Only the FCFS queue (a deque) has caller-controlled ordering; priority # queues order themselves and pair sync then relies on hold + equalize. if not isinstance(waiting, deque) or len(waiting) < 2: return requests = list(waiting) waiting.clear() seen: set[str] = set() result: list[Any] = [] for request in requests: rid = request.request_id if rid in seen: continue seen.add(rid) result.append(request) pair_id = scheduler._cfg_req_to_pair.get(rid) if pair_id is None: continue for partner_id in scheduler._cfg_pairs.get(pair_id, {}).values(): if partner_id != rid and partner_id not in seen: for candidate in requests: if candidate.request_id == partner_id: seen.add(partner_id) result.append(candidate) break waiting.extend(result) def _equalize_cfg_pair_progress(scheduler: Any, scheduler_output: Any) -> None: """Bring both members of every CFG pair to the same ``num_computed_tokens``. Prefix-cache hits or unequal chunked-prefill budget can leave one member ahead after a schedule step; reduce the faster member's allocation so both land on the same position. Over-allocated KV blocks remain and are consumed next step, so nothing leaks. """ for roles in scheduler._cfg_pairs.values(): 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 = all( req.num_computed_tokens - target < sched for req, sched in ((cond_req, cond_sched), (uncond_req, uncond_sched)) ) 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 _drop_split_pairs(scheduler: Any) -> None: """Deregister a desynced CFG pair (one member preempted / left behind). Manual re-preemption surgery is unsafe against the upstream scheduler. Instead, a pair that split anyway (KV-pressure preemption of exactly one member, or unrecoverable progress skew) is dropped from the registry: the cond request continues UNGUIDED (its logits are never blended again) rather than being re-paired against a desynced partner, and the companion's discarded output no longer matters. """ if not scheduler._cfg_pairs: return running_ids = {req.request_id for req in scheduler.running} to_drop: list[tuple[str, str, str, 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: # Split only counts once decoding started: a lone member that has # computed tokens while its partner was pushed back is desynced. runner = scheduler.requests.get(cond_id if cond_running else uncond_id) if runner is not None and runner.num_computed_tokens > 0: to_drop.append((pair_id, cond_id, uncond_id, "one member preempted")) 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: logger.warning( "CFG pair progress mismatch (equalize will retry): %s@%d vs %s@%d", cond_id, cond_req.num_computed_tokens, uncond_id, uncond_req.num_computed_tokens, ) for pair_id, cond_id, uncond_id, reason in to_drop: scheduler._cfg_pairs.pop(pair_id, None) scheduler._cfg_req_to_pair.pop(cond_id, None) scheduler._cfg_req_to_pair.pop(uncond_id, None) logger.warning( "CFG pair %s split (%s); releasing cond=%s unguided, uncond=%s output is discarded", pair_id, reason, cond_id, uncond_id, ) _patches_applied = False def cfg_patches_applied() -> bool: """True once :func:`apply_cfg_patches` has patched the scheduler in this process.""" return _patches_applied def _patch_scheduler(scheduler_cls: type | None = None) -> None: """Wrap the scheduler methods that keep a CFG pair step-locked. ``scheduler_cls`` defaults to vLLM's v1 ``Scheduler``; accepting a class keeps the wrappers exercisable without building a real engine. """ if scheduler_cls is None: from vllm.v1.core.sched.scheduler import Scheduler scheduler_cls = Scheduler orig_init = scheduler_cls.__init__ orig_add = scheduler_cls.add_request orig_finish = scheduler_cls.finish_requests orig_schedule = scheduler_cls.schedule def _init(self, *args: Any, **kwargs: Any) -> None: 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: Any) -> None: orig_add(self, request) pair_id = _cfg_extra(request, "cfg_pair_id") role = _cfg_extra(request, "cfg_role") if role and not pair_id: # Models that cannot inject a pair id from the serving layer send # the role alone and let the request id carry the pairing. It has # to be the user-facing id: the engine renames every request to an # internal UUID on admission, and the companion suffix survives # only on ``external_req_id``. Deriving from the UUID would give # the two members unrelated pair ids, so they would never be held # for each other and could be scheduled apart. external_id = getattr(request, "external_req_id", None) or request.request_id pair_id = _derive_pair_id(str(external_id), role) 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 # An asymmetric prefix-cache hit (the cond prompt shares cached # branches; the null prompt does not) can advance one member past # what post-schedule equalization can claw back. Skip prefix-cache # reads for CFG requests entirely. if hasattr(request, "skip_reading_prefix_cache"): request.skip_reading_prefix_cache = True def _finish(self, request_ids: Any, finished_status: Any) -> Any: 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) # Finish both members together: a surviving partner would decode # unpaired garbage (its logits are never blended again). 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 rid in self._cfg_pairs.get(pair_id, {}).values(): 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, *args: Any, **kwargs: Any) -> Any: if not self._cfg_pairs: return orig_schedule(self, *args, **kwargs) _reorder_waiting_for_cfg(self) held = _hold_incomplete_pairs(self) # Chunked prefill can split one pair member's prompt across steps # while the other fits in a single step; halving the long-prefill # threshold makes both chunk, and equalize aligns the remainder. scheduler_config = getattr(self, "scheduler_config", None) orig_threshold = getattr(scheduler_config, "long_prefill_token_threshold", None) if scheduler_config is not None and orig_threshold is not None: scheduler_config.long_prefill_token_threshold = self.max_num_scheduled_tokens // 2 try: scheduler_output = orig_schedule(self, *args, **kwargs) finally: if scheduler_config is not None and orig_threshold is not None: scheduler_config.long_prefill_token_threshold = orig_threshold _release_held(held) _equalize_cfg_pair_progress(self, scheduler_output) _drop_split_pairs(self) return scheduler_output _schedule._cfg_pairing_patched = True # type: ignore[attr-defined] # Legacy marker name kept for callers written against the Audex-only patch. _schedule._audex_cfg_patched = True # type: ignore[attr-defined] scheduler_cls.__init__ = _init scheduler_cls.add_request = _add scheduler_cls.finish_requests = _finish scheduler_cls.schedule = _schedule def apply_cfg_patches() -> None: """Make the vLLM v1 scheduler CFG-pair-aware. Idempotent, per process. Must run in the engine-core process before ``Scheduler`` is constructed (the ``__init__`` wrapper installs the pair registry). """ global _patches_applied if _patches_applied: return _patches_applied = True logger.info("Applying CFG pairing scheduler patches to vLLM v1") _patch_scheduler()