vllm-omni/vllm_omni/patch.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

import importlib
import logging
import os
import sys
from functools import cached_property

import torch
from aenum import extend_enum
from vllm.config import ModelConfig as _OriginalModelConfig
from vllm.config import VllmConfig as _OriginalVllmConfig
from vllm.inputs import TokensPrompt as _OriginalTokensPrompt
from vllm.model_executor.layers.rotary_embedding import (
    MRotaryEmbedding as _OriginalMRotaryEmbedding,
)
from vllm.v1.engine import EngineCoreOutput as _OriginalEngineCoreOutput
from vllm.v1.engine import EngineCoreOutputs as _OriginalEngineCoreOutputs
from vllm.v1.engine import EngineCoreRequest as _OriginalEngineCoreRequest
from vllm.v1.request import Request as _OriginalRequest
from vllm.v1.request import RequestStatus
from vllm.v1.request import StreamingUpdate as _OriginalStreamingUpdate

import vllm_omni.logger  # noqa: F401
from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs, OmniEngineCoreRequest
from vllm_omni.inputs.data import OmniTokensPrompt
from vllm_omni.model_executor.layers.rotary_embedding import OmniMRotaryEmbedding
from vllm_omni.request import OmniRequest, OmniStreamingUpdate

_PATCH_LOGGER = logging.getLogger("vllm_omni.patch")

# Omni's deploy config selects a runner per stage. vLLM 0.31's config property
# otherwise defaults to v2 independently of the runner selected by our workers,
# making shared consumers such as InputProcessor choose the wrong API.
_runner_property = _OriginalVllmConfig.use_v2_model_runner
if not getattr(_runner_property.fget, "_omni_stage_runner", False):

    def _patched_use_v2_model_runner(self, _upstream=_runner_property.fget):
        # Lazy import keeps the package's patch/config initialization acyclic.
        from vllm_omni.config.model import OmniModelConfig

        if isinstance(self.model_config, OmniModelConfig):
            return self.model_config.use_v2_model_runner
        return _upstream(self)

    _patched_use_v2_model_runner._omni_stage_runner = True
    _OriginalVllmConfig.use_v2_model_runner = property(_patched_use_v2_model_runner)

# =============================================================================
# Patch ModelConfig.is_mm_prefix_lm to support omni-specific models
# =============================================================================
# WHY: HunyuanImage-3.0 requires bidirectional attention for image tokens
# (cond_token_attn_type: "joint_full" in config.json). vLLM gates this on
# is_mm_prefix_lm, which checks an internal MM_PREFIX_LM_MODELS list that
# does not include "hunyuan_image_3_moe" (the upstream HF model_type).
#
# WHY NOT model-level: is_mm_prefix_lm is checked in vLLM core (scheduler,
# attention backend selection) before model code runs — no model-level hook.
#
# SCOPE: Only affects model_type in _OMNI_MM_PREFIX_LM_MODELS (currently
# just "hunyuan_image_3_moe"). All other models fall through to the
# original vLLM implementation unchanged.
#
# FRAGILITY: Relies on is_mm_prefix_lm being a cached_property on
# ModelConfig. The __dict__ access + __set_name__ dance works around a
# pydantic dataclass issue in vllm 0.19.0+. If vLLM changes
# is_mm_prefix_lm to a regular method or removes it, this will break.
#
# TODO: Upstream a configurable MM_PREFIX_LM_MODELS or a model_config flag
# so this patch can be removed.
_OMNI_MM_PREFIX_LM_MODELS = ("hunyuan_image_3_moe",)
# Access via __dict__ to avoid triggering cached_property.__get__ which fails
# with "Cannot use cached_property instance without calling __set_name__" in
# pydantic dataclasses (vllm 0.19.0+).
_cp = _OriginalModelConfig.__dict__["is_mm_prefix_lm"]
_original_is_mm_prefix_lm = _cp.func if hasattr(_cp, "func") else _cp.fget


def _patched_is_mm_prefix_lm(self):
    if _original_is_mm_prefix_lm(self):
        return True
    model_type = getattr(self.hf_config, "model_type", "")
    return model_type in _OMNI_MM_PREFIX_LM_MODELS


_patched_cp = cached_property(_patched_is_mm_prefix_lm)
_patched_cp.__set_name__(_OriginalModelConfig, "is_mm_prefix_lm")
_OriginalModelConfig.is_mm_prefix_lm = _patched_cp

# Sanity check: verify the patch is active. If vLLM changes the descriptor
# type or __set_name__ semantics, this will fail loudly at import time
# rather than silently falling back to unpatched behavior.
_installed = _OriginalModelConfig.__dict__.get("is_mm_prefix_lm")
assert _installed is _patched_cp, (
    "is_mm_prefix_lm patch failed to install — bidirectional attention "
    "for HunyuanImage3 will not work. Check vLLM ModelConfig changes."
)

# =============================================================================
# Patch ModelOptNvFp4LinearMethod.process_weights_after_loading to clamp NaN
# bytes in per-block FP8 weight_scale tensors at load time.
# =============================================================================
# WHY: ModelOpt 0.44's float32 -> torch.float8_e4m3fn cast of per-block weight
# scales occasionally emits NaN bytes (E4M3 encoding 0x7F / 0xFF) when the
# pre-cast scale rounds above the FP8 E4M3 max of 448 after the global-scale
# division. Any single NaN byte in weight_scale propagates through the
# FlashInfer FP4 GEMM and the served model collapses to `!!!!`. NVFP4 W4A4
# Qwen3-Omni checkpoints published by the community — exported with stock
# ModelOpt 0.44 — currently fail this way the moment they're served.
#
# WHY here (vllm-omni, not upstream): the root cause sits in ModelOpt; the
# analogous fix in vLLM's own PWAL is also pending. Both are planned as
# follow-up PRs. Until those land in vllm-omni's pinned vLLM, this load-time
# clamp keeps the user-visible failure from looking like a vllm-omni
# inference bug. Newly calibrated, clean checkpoints pay no runtime cost
# (the clamp is a no-op when no NaN bytes are present).
#
# SCOPE: the W4A4 NVFP4 linear PWAL only. NvFp4FusedMoE / NvFp4W4A16 /
# CompressedTensors / Quark NVFP4 paths are not covered.
#
# SELF-EXTINGUISH (structural): vLLM #49381 redesigned the ModelOpt linear
# methods around one generic `ModelOptLinearMethod` built by
# `build_linear_method()`, and removed the per-format classes
# (`ModelOptNvFp4LinearMethod`, `ModelOptFp8LinearMethod`, ...) together with
# the `LinearMethodCls` attributes they were reached through. This block
# therefore resolves its target dynamically and self-extinguishes when only
# the generic class exists, because that PWAL now REJECTS any NaN
# weight_scale (#52501: `KNvfp4Static` creates `weight_scale` as
# `torch.full(shape, nan)` and raises "... was never loaded (still NaN)" on
# any surviving NaN). A NaN byte in weight_scale can come from that
# unloaded-scale sentinel or from ModelOpt 0.44's FP32→FP8 E4M3 cast
# overflow, and the two are byte-identical; a fused projection's
# weight_scale is also written slice-by-slice by the weight loader, so
# "some values are finite" does not prove the tensor was fully loaded.
# Clamping would silently serve a partially loaded model, so the override is
# retired there instead — a corrupt checkpoint gets a loud load-time
# RuntimeError instead of the `!!!!` decode-time collapse.
#
# SELF-EXTINGUISH (heuristic, legacy pins only): `_already_patched_upstream`
# heuristically detects when vLLM's own PWAL contains an in-place
# `masked_fill_` against `weight_scale` / `isnan` — the structure the
# upstream fix was expected to take. NOTE: the heuristic only matches
# "vLLM PWAL clamps NaN in-place with `masked_fill_`"; if an upstream fix
# lands as `nan_to_num_` or as a clamp before the FP32→FP8 cast, the check
# won't fire and this override stays active. The override is idempotent so
# the overlap is a warning log, not a correctness issue.
#
# ORDERING: the clamp must run BEFORE the original PWAL. The non-Blackwell
# Marlin fallback (sm_<100) casts weight_scale FP8 -> bf16/fp16 and permutes
# inside its kernel PWAL; an after-PWAL clamp would either trip the
# byte-view shape assertion (FP8 1B/elem -> bf16 2B/elem) or operate on
# already-transformed bytes where NaN has propagated through permute. On
# Blackwell (Cutlass/TRTLLM/cuDNN) the kernel PWAL only swizzles/shuffles
# FP8 bytes, so clamp-before is equivalent to clamp-after there.


def _clamp_nvfp4_weight_scale_nans(layer) -> int:
    """Scan ``layer.weight_scale`` for NaN bytes and clamp them to the FP8
    E4M3 max byte (0x7E) in place. Returns the number of bytes clamped.

    Scoped to NVFP4 ``float8_e4m3fn`` weight_scale only — a future MXFP4
    weight-scale clamp would need its own helper (MXFP4 uses
    ``float8_e8m0fnu`` for scales, and the byte encoding of NaN differs).

    Exposed at module scope so tests can exercise the clamp directly
    without monkey-patching + reloading the parent ``vllm_omni.patch``
    module.
    """
    # Defensive: a stray subclass without weight_scale would crash inside
    # the override; just bail out.
    weight_scale = getattr(layer, "weight_scale", None)
    if weight_scale is None:
        return 0
    # Re-entry safety: this clamp expects raw on-disk FP8 weight_scale (1
    # byte/elem). If we are ever invoked after the original PWAL ran (e.g.
    # a caller re-issues process_weights_after_loading on the same layer),
    # weight_scale may have been cast to bf16/fp16 by the Marlin path, in
    # which case the uint8 byte view below would have a different shape
    # than the NaN mask. Bail out cleanly rather than trip the assert.
    if weight_scale.dtype != torch.float8_e4m3fn:
        return 0
    # `.view(torch.uint8)` below requires a contiguous tensor and raises on a
    # non-contiguous one. A checkpoint whose per-block scale tensor isn't
    # contiguous would otherwise crash here. Materialize a contiguous copy and
    # write it back to the Parameter so the in-place byte clamp lands on the
    # tensor the kernel will actually read (a bare `.contiguous()` copy would
    # be discarded and the original left un-clamped).
    weight_scale_data = weight_scale.data
    if not weight_scale_data.is_contiguous():
        weight_scale_data = weight_scale_data.contiguous()
        weight_scale.data = weight_scale_data
    nan_mask = torch.isnan(weight_scale_data)
    if not nan_mask.any():
        return 0
    # PyTorch does not implement masked_fill_ for float8_e4m3fn, so view
    # the storage as uint8 and write the FP8 max byte (0x7E) directly.
    # float8_e4m3fn is 1 byte per element so the byte view is shape-
    # for-shape with the original tensor; pin that here so a future
    # packed scale format would fail loudly rather than silently corrupt
    # unrelated bytes.
    byte_view = weight_scale_data.view(torch.uint8)
    if byte_view.shape != nan_mask.shape:
        raise RuntimeError(
            f"NVFP4 weight_scale uint8 view shape {byte_view.shape} != NaN mask shape "
            f"{nan_mask.shape}; per-block scale layout changed — recheck the clamp."
        )
    fp8_max_byte = (
        torch.tensor(torch.finfo(torch.float8_e4m3fn).max, dtype=torch.float8_e4m3fn).view(torch.uint8).item()
    )
    byte_view.masked_fill_(nan_mask, fp8_max_byte)
    n = int(nan_mask.sum().item())
    _PATCH_LOGGER.warning("Clamped %d NaN entries in NVFP4 per-block weight_scale.", n)
    return n


# Module-level defaults so downstream code (and tests) can import these names
# without guarding for the import-failure / escape-hatch branches below.
# `_already_patched_upstream` = upstream's own PWAL already handles NaN
#                               weight_scale by itself (its own clamp, or — on
#                               the redesigned ModelOpt path — a hard
#                               rejection), so we deliberately install nothing.
# `_clamp_installed`         = our wrapper was installed on the upstream class.
# These are independent: the env-var escape hatch and the import-failure path
# both leave the wrapper uninstalled WITHOUT upstream handling the case, so the
# right check for "we own NaN-clamp behavior" is `_clamp_installed`.
_already_patched_upstream = False
_clamp_installed = False

# Resolve the clamp target dynamically. A literal
# `from ...modelopt import ModelOptNvFp4LinearMethod` is not an option any more:
# the class was removed upstream (#49381), so the import is both a hard
# ImportError and a static-check failure, and dynamic resolution is also what
# lets us tell "no target" apart from "target replaced by the generic method".
_MODELOPT_MODULE = "vllm.model_executor.layers.quantization.modelopt"
_LEGACY_NVFP4_LINEAR_METHOD = "ModelOptNvFp4LinearMethod"
_GENERIC_LINEAR_METHOD = "ModelOptLinearMethod"

try:
    # Escape hatch — set VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP=1 to skip installing
    # the patch (e.g. to confirm a `!!!!` failure is the NaN-byte case).
    # The escape-hatch deliberately raises ImportError so the not-installed
    # warning below logs through the same path a real ImportError would, and it
    # short-circuits BEFORE any upstream probing so both module flags stay
    # False on this path.
    # Use the repo-wide bool-env idiom so values like `0`, `false`, `no`,
    # `off` correctly mean "do not skip" rather than tripping naive
    # truthiness on the non-empty string.
    if os.environ.get("VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP", "").lower() in ("1", "true", "yes", "on"):
        raise ImportError("VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP is set; skipping NaN-clamp install")
    _modelopt = importlib.import_module(_MODELOPT_MODULE)
    _legacy_nvfp4_linear_method = getattr(_modelopt, _LEGACY_NVFP4_LINEAR_METHOD, None)
    _generic_linear_method = getattr(_modelopt, _GENERIC_LINEAR_METHOD, None)
except ImportError as _nan_clamp_import_err:
    _PATCH_LOGGER.warning(
        "NVFP4 weight_scale NaN-clamp patch could NOT install: %s. NVFP4 W4A4 "
        "checkpoints with NaN bytes in per-block weight_scale will serve `!!!!`.",
        _nan_clamp_import_err,
    )
    _legacy_nvfp4_linear_method = None
    _generic_linear_method = None
else:
    if _legacy_nvfp4_linear_method is None and _generic_linear_method is not None:
        # Redesigned upstream (#49381): only the generic ModelOptLinearMethod
        # exists, and its PWAL rejects any NaN weight_scale (#52501). Retire the
        # override rather than mask that check — see the SELF-EXTINGUISH
        # (structural) note above for why the two NaN origins cannot be told
        # apart. `_already_patched_upstream` records that upstream handles the
        # case itself, so an absent `_clamp_installed` is the expected state.
        _already_patched_upstream = True
        _PATCH_LOGGER.info(
            "NVFP4 W4A4 weight_scale NaN-clamp: skipped — upstream serves ModelOpt "
            "linears through the generic %s and rejects any NaN weight_scale "
            "(unloaded-scale sentinel), so a load-time error replaces the "
            "`!!!!` decode-time collapse.",
            _GENERIC_LINEAR_METHOD,
        )
    elif _legacy_nvfp4_linear_method is None:
        # Unrecognised pin: neither the legacy per-format class nor the generic
        # replacement is present. Nothing to install and nothing to assume.
        _PATCH_LOGGER.warning(
            "NVFP4 weight_scale NaN-clamp patch could NOT install: neither %s nor %s is present in %s.",
            _LEGACY_NVFP4_LINEAR_METHOD,
            _GENERIC_LINEAR_METHOD,
            _MODELOPT_MODULE,
        )

if _legacy_nvfp4_linear_method is not None:
    _OriginalModelOptNvFp4LinearMethod = _legacy_nvfp4_linear_method
    _current_nvfp4_pwal = _OriginalModelOptNvFp4LinearMethod.process_weights_after_loading
    # Reload idempotency: on a module reload (importlib.reload in a test, or a
    # second import path) the class attribute already holds OUR wrapper, so
    # capturing it as the "original" and re-wrapping would nest the clamp and
    # run it twice per load. Recover the genuine upstream method from the
    # sentinel we stamp below; the heuristic + a fresh single-level wrapper are
    # then computed against the real upstream, not against our own wrapper.
    _original_nvfp4_pwal = getattr(_current_nvfp4_pwal, "_vllm_omni_wrapped_pwal", _current_nvfp4_pwal)
    _upstream_pwal_names = set(_original_nvfp4_pwal.__code__.co_names or ())
    # Require ALL three names — `weight_scale` alone is too loose because the
    # current upstream PWAL already references `weight_scale_2` (close prefix
    # collisions aside, future PWALs may legitimately reference `weight_scale`
    # without a NaN clamp). `masked_fill_` + `isnan` + `weight_scale` together
    # are the structural signature of an in-place NaN clamp on weight_scale.
    # KNOWN BLIND SPOTS (false-negatives, safe direction — we install a
    # redundant but idempotent clamp): upstream using
    # `getattr(layer, "weight_scale")` (string lands in co_consts, not
    # co_names), or factoring the clamp into a helper called from PWAL
    # (none of the three names appear in the top-level co_names). Revisit
    # this heuristic when the upstream PR is actually filed.
    _already_patched_upstream = all(n in _upstream_pwal_names for n in ("masked_fill_", "isnan", "weight_scale"))

    def _patched_nvfp4_pwal(self, layer, *args, **kwargs):
        # Clamp BEFORE the original PWAL — see ORDERING note above.
        _clamp_nvfp4_weight_scale_nans(layer)
        _original_nvfp4_pwal(self, layer, *args, **kwargs)

    # Sentinel so a later reload of this module recovers the genuine upstream
    # method above instead of wrapping this wrapper (see reload-idempotency note).
    _patched_nvfp4_pwal._vllm_omni_wrapped_pwal = _original_nvfp4_pwal

    if not _already_patched_upstream:
        _OriginalModelOptNvFp4LinearMethod.process_weights_after_loading = _patched_nvfp4_pwal
        # Fail loudly if install dropped silently (e.g. another plugin
        # patched the same class) rather than degrade to `!!!!` at decode.
        # Use raise, not assert: asserts are compiled out under `python -O` /
        # PYTHONOPTIMIZE, which would silently disable this guard in exactly
        # the optimized runs where a conflicting plugin must still be caught.
        if _OriginalModelOptNvFp4LinearMethod.process_weights_after_loading is not _patched_nvfp4_pwal:
            raise RuntimeError("NVFP4 weight_scale NaN-clamp install failed — check for conflicting plugins.")
        _clamp_installed = True

    _PATCH_LOGGER.info(
        "NVFP4 W4A4 weight_scale NaN-clamp: %s.",
        "skipped (upstream already patched)" if _already_patched_upstream else "installed",
    )

# =============================================================================
# Patch GlmImageTextConfig to expose mrope_section in rope_parameters
# =============================================================================
# GLM-Image uses M-RoPE with mrope_section: [8, 12, 12], but transformers'
# implementation doesn't expose it in rope_parameters. vLLM's uses_mrope
# detection relies on "mrope_section" being present in rope_parameters.
# This patch ensures proper M-RoPE detection for GLM-Image.
try:
    from transformers.models.glm_image.configuration_glm_image import GlmImageTextConfig

    _original_glm_image_text_config_init = GlmImageTextConfig.__init__

    def _patched_glm_image_text_config_init(self, *args, **kwargs):
        _original_glm_image_text_config_init(self, *args, **kwargs)
        # Ensure rope_parameters exists and contains mrope_section
        if self.rope_parameters is None:
            self.rope_parameters = {}
        if isinstance(self.rope_parameters, dict) and "mrope_section" not in self.rope_parameters:
            # GLM-Image uses mrope_section: [8, 12, 12] for T/H/W dimensions
            self.rope_parameters["mrope_section"] = [8, 12, 12]

    GlmImageTextConfig.__init__ = _patched_glm_image_text_config_init
except ImportError:
    # GlmImageTextConfig not available, skip patching
    pass

# Extend RequestStatus enum with omni-specific statuses
if not hasattr(RequestStatus, "WAITING_FOR_CHUNK"):
    # The value - 1 is intentionally chosen to ensure it is treated
    # as a non-finished state and remains compatible with existing comparisons.
    extend_enum(RequestStatus, "WAITING_FOR_CHUNK", -1)

if not hasattr(RequestStatus, "WAITING_FOR_INPUT"):
    # Full-payload stage handoff uses a distinct waiting state so the
    # scheduler can restore the request once non-stage-0 inputs arrive.
    extend_enum(RequestStatus, "WAITING_FOR_INPUT", -2)

# Snapshot sys.modules: `hasattr` below can trigger lazy submodule imports
# (e.g. transformers' `_LazyModule.__getattr__`), which mutate sys.modules
# during iteration and raise `dictionary changed size during iteration`.
for module_name, module in list(sys.modules.items()):
    # only do patch on module of vllm, pass others
    if "vllm" not in module_name:
        continue
    if hasattr(module, "EngineCoreOutput") and module.EngineCoreOutput == _OriginalEngineCoreOutput:
        module.EngineCoreOutput = OmniEngineCoreOutput
    if hasattr(module, "EngineCoreOutputs") and module.EngineCoreOutputs == _OriginalEngineCoreOutputs:
        module.EngineCoreOutputs = OmniEngineCoreOutputs
    if hasattr(module, "TokensPrompt") and module.TokensPrompt == _OriginalTokensPrompt:
        module.TokensPrompt = OmniTokensPrompt
    if hasattr(module, "MRotaryEmbedding") and module.MRotaryEmbedding == _OriginalMRotaryEmbedding:
        module.MRotaryEmbedding = OmniMRotaryEmbedding
    if hasattr(module, "Request") and module.Request == _OriginalRequest:
        module.Request = OmniRequest
    if hasattr(module, "StreamingUpdate") and module.StreamingUpdate == _OriginalStreamingUpdate:
        module.StreamingUpdate = OmniStreamingUpdate
    if hasattr(module, "EngineCoreRequest") and module.EngineCoreRequest == _OriginalEngineCoreRequest:
        module.EngineCoreRequest = OmniEngineCoreRequest


# Patch: add qwen3_omni_moe to vllm's chat template fallback registry.
# Qwen/Qwen3-Omni-30B-A3B-Instruct stores its chat_template in a standalone
# chat_template.json (not in tokenizer_config.json).  transformers < 5.9.0
# does not load this file, so the tokenizer has no chat_template attribute.
# vllm's resolve_chat_template falls back to MODEL_TYPE_TO_CHAT_TEMPLATE
# which has "qwen" but not "qwen3_omni_moe".  Register the same fallback.
def _patch_chat_template_registry():
    try:
        from vllm.transformers_utils.chat_templates.registry import (
            _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK,
            CHAT_TEMPLATES_DIR,
        )

        if "qwen3_omni_moe" not in _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK:
            _MODEL_TYPE_TO_CHAT_TEMPLATE_FALLBACK["qwen3_omni_moe"] = lambda _: (
                CHAT_TEMPLATES_DIR / "template_chatml.jinja"
            )
    except ImportError:
        pass


_patch_chat_template_registry()


def _patch_scaled_mm_fp8_contiguous_activation():
    """Support batched diffusion activations on the ModelOpt FP8 (ScaledMM) path.

    The FP8 ScaledMM linear flattens its activation with ``x.view(-1, ...)``, which
    needs a contiguous tensor. Under step-execution batching (``--max-num-seqs > 1``)
    the sequence-packed diffusion activations can be non-contiguous, so we make the
    activation contiguous before the GEMM (no-op when it already is). Mixed FP8/NVFP4
    routes through the CUTLASS NVFP4 path and is unaffected.
    """
    try:
        from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
            ScaledMMLinearKernel,
        )
    except ImportError:
        return

    _original_apply_weights = ScaledMMLinearKernel.apply_weights

    def _patched_apply_weights(self, layer, x, bias=None):
        if not x.is_contiguous():
            x = x.contiguous()
        return _original_apply_weights(self, layer, x, bias)

    ScaledMMLinearKernel.apply_weights = _patched_apply_weights


_patch_scaled_mm_fp8_contiguous_activation()


def _patch_flashinfer_fp8_scaled_mm_output_shape():
    """Restore the N-D output shape for the FlashInfer FP8 ScaledMM kernel.

    ``FlashInferFP8ScaledMMLinearKernel.apply_scaled_mm`` returns the raw 2-D
    GEMM result and ignores ``output_shape``, unlike the CUTLASS / PyTorch
    ScaledMM kernels which reshape to it. A 3-D activation ``(B, S, D)`` thus
    collapses to ``(B*S, D)``, breaking diffusion DiTs that reshape the linear
    output by absolute dim (e.g. Wan2.2 ``qkv.unflatten(2, ...)``) with
    ``IndexError: Dimension out of range``. It only bites >2-D inputs, so LLM
    (token-flattened, 2-D) paths are unaffected.

    Carried here because the upstream fix may not have landed yet; this override
    becomes a harmless no-op once vLLM honors ``output_shape`` itself.
    """
    try:
        from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import (
            FlashInferFP8ScaledMMLinearKernel,
        )
    except ImportError:
        return

    _original_apply_scaled_mm = FlashInferFP8ScaledMMLinearKernel.apply_scaled_mm
    if getattr(_original_apply_scaled_mm, "_omni_output_shape_patched", False):
        return

    def _patched_apply_scaled_mm(self, *, A, B, out_dtype, As, Bs, bias, output_shape):  # noqa: N803
        out = _original_apply_scaled_mm(
            self, A=A, B=B, out_dtype=out_dtype, As=As, Bs=Bs, bias=bias, output_shape=output_shape
        )
        if tuple(out.shape) != tuple(output_shape):
            out = out.view(*output_shape)
        return out

    _patched_apply_scaled_mm._omni_output_shape_patched = True
    FlashInferFP8ScaledMMLinearKernel.apply_scaled_mm = _patched_apply_scaled_mm


_patch_flashinfer_fp8_scaled_mm_output_shape()


def _patch_fp8_use_quack_fused_bias():
    try:
        from vllm_omni.quantization.quack_fp8 import install_quack_fp8_patch

        install_quack_fp8_patch()
    except Exception:  # noqa: BLE001
        pass


_patch_fp8_use_quack_fused_bias()


# =============================================================================
# Patch torch inductor: prove factorable symbolic divisibility (CantSplit)
# =============================================================================
# WHY: torch 2.13's SizeVarAllocator.statically_known_multiple_of proves
# symbolic divisibility via torch's own Mod (torch.utils._sympy.functions),
# which deliberately stays unevaluated to bound compile time on wide
# expressions. torch <= 2.11 used python `%` (sympy.Mod), whose eval factors
# common terms. As a result an expression like
#   15360*s31 + 15360*s87  vs  s31 + s87
# — FLUX's fp8 graph with two dynamic sequence dims — is no longer provably
# divisible, and inductor raises CantSplit on the first compiled forward,
# killing the engine core during the warmup dummy run (nightly builds
# 2953/2954, Diffusion Quantization Test).
#
# SCOPE: wraps statically_known_multiple_of. The wrapper can only ADD True
# results, and only when sympy.cancel(numerator/denominator) yields a
# polynomial with integer coefficients — then numerator == denominator * q
# identically, so divisibility holds for every symbol assignment; this is
# unconditionally sound. Everything else defers to the original result.
# Guards: polynomial inputs only (inductor's FloorDiv/ModularIndexing atoms
# never reach sympy.cancel) and <= 20 free symbols (mirrors upstream's own
# cost cap).
#
# FRAGILITY / REMOVE WHEN: self-extinguishes at install time by probing
# whether the unpatched method already proves the canonical factorable case
# (as torch <= 2.11 does). When upstream restores factor-aware proving, the
# probe passes and nothing is patched.
def _provably_factorable_multiple(numerator, denominator) -> bool:
    """True only when numerator/denominator cancels to an integer polynomial."""
    try:
        import sympy

        num = sympy.sympify(numerator)
        den = sympy.sympify(denominator)
        if den.is_zero:
            return False
        symbols = num.free_symbols | den.free_symbols
        if len(symbols) > 20:
            return False
        if not (num.is_polynomial(*symbols) and den.is_polynomial(*symbols)):
            return False
        quotient = sympy.cancel(num / den)
        _, q_den = sympy.fraction(sympy.together(quotient))
        if q_den != 1:
            return False
        if quotient.is_Integer:
            return True
        poly = sympy.Poly(quotient, *sorted(quotient.free_symbols, key=str))
        return all(coeff.is_integer for coeff in poly.coeffs())
    except Exception:  # noqa: BLE001 - a proof failure must never break compile
        return False


def _patch_inductor_factorable_divisibility():
    try:
        import sympy
        from torch._inductor.sizevars import SizeVarAllocator
    except ImportError:
        return

    original = SizeVarAllocator.statically_known_multiple_of
    if getattr(original, "_vllm_omni_factorable_divisibility", False):
        return

    try:
        _a, _b = sympy.symbols("_omni_probe_a _omni_probe_b", positive=True, integer=True)
        if original(SizeVarAllocator(), 7 * _a + 7 * _b, _a + _b):
            _PATCH_LOGGER.info("inductor factorable-divisibility patch: skipped (upstream already proves it).")
            return
    except Exception:  # noqa: BLE001
        # Probe broke (constructor drift etc.) — install anyway; the wrapper
        # never subtracts results from the original.
        pass

    def statically_known_multiple_of(self, numerator, denominator):
        if original(self, numerator, denominator):
            return True
        return _provably_factorable_multiple(numerator, denominator)

    statically_known_multiple_of._vllm_omni_factorable_divisibility = True
    SizeVarAllocator.statically_known_multiple_of = statically_known_multiple_of
    _PATCH_LOGGER.info("inductor factorable-divisibility patch: installed.")


_patch_inductor_factorable_divisibility()


# =============================================================================
# Patch CuMemAllocator._python_free_callback to fix CUDA double-free on shutdown
# =============================================================================
# WHY: CuMemAllocator._python_free_callback guards the asleep-entry double-free
# skip with ``if data.is_asleep and current_platform.is_rocm():`` — only ROCm
# gets the safe empty-handle return.  On CUDA, the callback falls through and
# returns the original handle, causing cuMemRelease on already-freed memory
# (CUDA_ERROR_INVALID_VALUE) during EngineCore subprocess atexit cleanup.
#
# This happens because ``sleep()`` calls ``unmap_and_release()`` on ALL
# platforms, then sets ``is_asleep = True``.  When the atexit handler
# (``_shutdown_singleton`` -> ``release_pools()`` -> GC) triggers the free
# callback, the asleep entries must return an empty chunk list so the C
# extension skips ``cuMemRelease`` — exactly the same logic already used by
# the ROCm guard.
#
# The fix removes the ``current_platform.is_rocm()`` condition so the guard
# applies to CUDA (and any future platform that implements cumem).
#
# FRAGILITY: Relies on ``_python_free_callback`` being a regular method
# (not a slot or C extension).  If CuMemAllocator is rewritten in C/Cython,
# this monkey-patch will silently become a no-op.
def _patch_cumem_free_callback_cuda() -> None:
    try:
        from vllm.device_allocator.cumem import CuMemAllocator
    except ImportError:
        _PATCH_LOGGER.debug("[cumem-cuda] CuMemAllocator not available; skipping patch")
        return

    _original_free_callback = CuMemAllocator._python_free_callback

    if getattr(_original_free_callback, "_omni_cumem_cuda_patched", False):
        return

    # The upstream bug: `_python_free_callback` only skips the double-free
    # for ROCm (line ~206).  We wrap the method to extend the guard to all
    # platforms.
    def _patched_free_callback(self, ptr: int) -> tuple:
        data = self.pointer_to_data.pop(ptr)
        if data.cpu_backup_tensor is not None:
            data.cpu_backup_tensor = None
        if data.is_asleep:
            # sleep() already called unmap_and_release() on this allocation.
            # Return an empty chunk list so the C extension skips
            # cuMemRelease, avoiding a double-free.  Same logic as the
            # existing ROCm guard, but applied to all platforms.
            device, size, d_mem, _ = data.handle
            result = (device, size, d_mem, [])
            _PATCH_LOGGER.debug(
                "[cumem-cuda] Free callback: asleep entry %s -> empty handle",
                ptr,
            )
            return result
        # Drain pending kernels before the C extension's cuMemUnmap.
        torch.accelerator.synchronize(data.handle[0])
        return data.handle

    _patched_free_callback._omni_cumem_cuda_patched = True
    CuMemAllocator._python_free_callback = _patched_free_callback
    _PATCH_LOGGER.info(
        "[cumem-cuda] CuMemAllocator._python_free_callback patched: asleep guard extended to all platforms."
    )


_patch_cumem_free_callback_cuda()


def _patch_batched_seeded_random_sample() -> None:
    """Draw per-request seeded sampling noise in one launch.

    vLLM's ``random_sample`` calls ``exponential_(generator=...)`` once per
    seeded request and step. The replacement reproduces those draws bit for
    bit (values and generator states) with one kernel; unsupported cases keep
    the original loop.
    """
    try:
        import inspect

        from vllm.v1.sample.ops import topk_topp_sampler as sampler_ops
    except ImportError:
        return
    original = getattr(sampler_ops, "random_sample", None)
    if original is None or getattr(original, "_omni_batched_seeded", False):
        return
    if list(inspect.signature(original).parameters) != ["probs", "generators", "use_fp64_gumbel"] or not all(
        hasattr(sampler_ops, name) for name in ("empty_exponential_noise_like", "sample_with_exponential_noise")
    ):
        _PATCH_LOGGER.debug("[seeded-sampling] random_sample signature changed; batching disabled")
        return

    def random_sample(
        probs: torch.Tensor,
        generators: dict[int, torch.Generator],
        use_fp64_gumbel: bool = False,
    ) -> torch.Tensor:
        q = sampler_ops.empty_exponential_noise_like(probs, use_fp64_gumbel)
        if len(generators) != probs.shape[0]:
            q.exponential_()
        if generators:
            from vllm_omni.utils.seeded_exponential import (
                batched_seeded_exponential_supported,
                fill_exponential_rows,
            )

            if batched_seeded_exponential_supported(q, generators):
                fill_exponential_rows(q, list(generators.values()), list(generators))
            else:
                for i, generator in generators.items():
                    q[i].exponential_(generator=generator)
        return sampler_ops.sample_with_exponential_noise(probs, q)

    random_sample._omni_batched_seeded = True  # type: ignore[attr-defined]
    random_sample.__wrapped__ = original  # type: ignore[attr-defined]
    sampler_ops.random_sample = random_sample


_patch_batched_seeded_random_sample()


def _patch_mrv2_stop_token_capacity() -> None:
    """Let Model Runner V2 enforce ``min_tokens`` over large stop sets.

    MRV2 keeps each request's stop token ids for ``min_tokens`` masking in a
    fixed-width GPU table capped at 128 ids and rejects larger sets. Speech
    LMs stop on every special id past the codebook (CosyVoice3: 200). The
    masking kernel sizes its block from the table widths, which other tables
    already raise to 1024, so widening this one only grows a small buffer.
    """
    try:
        from vllm.v1.worker.gpu.sample import logit_bias
    except ImportError:
        return
    logit_bias.MAX_NUM_STOP_TOKEN_IDS = max(
        logit_bias.MAX_NUM_STOP_TOKEN_IDS,
        logit_bias.MAX_NUM_ALLOWED_TOKEN_IDS,
    )


_patch_mrv2_stop_token_capacity()