inference_scripts_vllm/audiogen_scripts/vllm_cfg_patch.py

# 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()