From 2f64c712fdad324c73069d61e6b10fb041b018f2 Mon Sep 17 00:00:00 2001 From: Yihong Jin Date: Wed, 17 Jun 2026 19:01:46 -0700 Subject: [PATCH] [Quantization][Qwen3-Omni] Enable NVFP4 W4A4 serving on Blackwell (#4025) Signed-off-by: Yihong Jin --- docs/user_guide/quantization/modelopt.md | 43 +++ .../test_qwen3_omni_modelopt_nvfp4_w4a4.py | 100 +++++++ .../test_modelopt_nvfp4_nan_clamp.py | 279 ++++++++++++++++++ vllm_omni/patch.py | 183 ++++++++++++ 4 files changed, 605 insertions(+) create mode 100644 tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py create mode 100644 tests/model_executor/quantization/test_modelopt_nvfp4_nan_clamp.py diff --git a/docs/user_guide/quantization/modelopt.md b/docs/user_guide/quantization/modelopt.md index 66e8427bb..d8da4e1a9 100644 --- a/docs/user_guide/quantization/modelopt.md +++ b/docs/user_guide/quantization/modelopt.md @@ -89,11 +89,54 @@ and | Model | Scope | Status | |-------|-------|--------| | Qwen3-Omni | Thinker language-model stage | ModelOpt FP8 checkpoint path | +| Qwen3-Omni | Thinker language-model stage (W4A4 NVFP4) | Validated; see [Qwen3-Omni NVFP4 W4A4](#qwen3-omni-nvfp4-w4a4-thinker) below | | Qwen3-TTS | TTS language-model stage | Not validated | Audio encoder, vision encoder, talker, and code2wav stages stay in BF16 unless a model-specific guide documents otherwise. +#### Qwen3-Omni NVFP4 W4A4 (thinker) + +vLLM-Omni serves ModelOpt NVFP4 W4A4 quantizations of the +Qwen3-Omni-30B-A3B-Instruct thinker language model. The thinker text body +(attention + MoE experts) is quantized to NVFP4 with FP8 per-tensor input +scales; the audio encoder, vision encoder, talker, and code2wav stay in BF16. + +| Variant | HF checkpoint | Hardware | +|---------|---------------|----------| +| W4A4 NVFP4 (full thinker) | `YihongJin/Qwen3-Omni-30B-A3B-Instruct-NVFP4-W4A4-full-thinker-awqclip` | sm_100+ (Blackwell, FlashInfer FP4 GEMM) | + +Calibration uses ModelOpt `mtq.NVFP4_DEFAULT_CFG` with the `awq_clip` algorithm +on 1024 ultrachat samples chat-templated through the Qwen3-Omni tokenizer. +Excluded modules: `*audio_tower*`, `*visual*`, `*talker*`, `*code2wav*`, +`*lm_head*`, `*mlp.gate*`. See `scripts/nvfp4/calibrate.py` for the reference +recipe. + +!!! note "ModelOpt 0.44 NaN regression workaround" + ModelOpt 0.44's float32 -> FP8 E4M3 cast of per-block weight scales + occasionally emits literal NaN bytes (E4M3 encoding 0x7F / 0xFF) for + blocks whose pre-cast scale rounds above the FP8 max of 448 after the + global-scale division. A single NaN byte in any `weight_scale` + propagates through the FP4 GEMM and collapses the served model output + to `!!!!`. vLLM-Omni's `vllm_omni.patch` installs a defensive override + of `ModelOptNvFp4LinearMethod.process_weights_after_loading` that + clamps these bytes to the FP8 E4M3 max at load time. The override + self-extinguishes once vllm-omni's vllm pin moves to a release + containing the upstream fix. Set `VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP=1` + to disable the override for diagnostics. + +Serving: + +```bash +vllm serve YihongJin/Qwen3-Omni-30B-A3B-Instruct-NVFP4-W4A4-full-thinker-awqclip \ + --omni --port 8000 +``` + +> Do **not** pass `--enforce-eager` for production / benchmarks. CUDA graphs +> amortize launch overhead and unlock the FP4 throughput wins; with +> `--enforce-eager` set, W4A4 TPOT degrades ~10x relative to the CUDA-graph +> configuration. + ### Multi-Stage Diffusion Model ModelOpt checkpoints must be routed to the stage whose checkpoint contains the diff --git a/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py b/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py new file mode 100644 index 000000000..5d1b3dda7 --- /dev/null +++ b/tests/e2e/offline_inference/test_qwen3_omni_modelopt_nvfp4_w4a4.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""E2E tests for Qwen3-Omni ModelOpt NVFP4 W4A4 — Marlin fallback path. + +Loads the published full-thinker NVFP4 W4A4 checkpoint on H100 (sm_90), +where ``init_nvfp4_linear_kernel()`` selects the ``MarlinNvFp4LinearKernel`` +weight-only fallback. The Marlin kernel's PWAL casts ``weight_scale`` +FP8 -> bf16 and permutes; the patched +``ModelOptNvFp4LinearMethod.process_weights_after_loading`` in +``vllm_omni.patch`` therefore clamps NaN bytes BEFORE delegating to the +original PWAL. This test asserts the served response is coherent, i.e. +no `!!!!` token cascade caused by NaN bytes in per-block ``weight_scale``. + +Hardware coverage caveat: this DOES NOT exercise the native FlashInfer +FP4 GEMM (Cutlass / TRT-LLM / cuDNN) that is selected on Blackwell +sm_100+ — those kernel PWALs swizzle/shuffle FP8 bytes without dtype +casting, a different code path from Marlin. TODO(B200): add a companion +Blackwell smoke once a Blackwell runner is wired into ``tests.helpers.mark``. +""" + +import os +import re + +import pytest + +from tests.helpers.mark import hardware_test +from tests.helpers.stage_config import get_deploy_config_path, modify_stage_config + +pytestmark = [ + pytest.mark.full_model, + pytest.mark.omni, +] + +QUANTIZED_MODEL = os.environ.get( + "QWEN3_OMNI_W4A4_MODEL", + "YihongJin/Qwen3-Omni-30B-A3B-Instruct-NVFP4-W4A4-full-thinker-awqclip", +) + +_CI_DEPLOY = get_deploy_config_path("ci/qwen3_omni_moe.yaml") + + +# `!!!!` (4+ exclamation marks) is the signature degenerate output the +# NaN-clamp patch defends against. Any test that loads a W4A4 checkpoint +# must reject this. +_NAN_COLLAPSE_RE = re.compile(r"!{4,}") + + +@pytest.fixture(scope="module", autouse=True) +def _qwen3_omni_env(): + with pytest.MonkeyPatch.context() as mp: + mp.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn") + yield + + +def _stage_config(): + # enforce_eager=True here to keep CI memory + warm-up time bounded for the + # smoke test, matching the existing Qwen3-Omni e2e pattern (see + # test_qwen3_omni_autoround_w4a16_expansion.py). The throughput claims in + # the PR body are from explicit no-enforce-eager runs, not from this test. + # + # TODO: the production code path uses FULL_AND_PIECEWISE CUDA graphs. + # The NaN clamp itself is graph-mode-independent so this enforce-eager + # smoke is still a valid regression guard against `!!!!`, but a future + # CUDA-graph-specific regression would not be caught here. If/when CI + # gains enough VRAM headroom for a no-eager run, add a companion test. + return modify_stage_config( + _CI_DEPLOY, + updates={"stages": {0: {"enforce_eager": True}, 1: {"enforce_eager": True}}}, + ) + + +quant_params = [(QUANTIZED_MODEL, _stage_config())] + + +@hardware_test(res={"cuda": "H100"}, num_cards=2) +@pytest.mark.parametrize("omni_runner", quant_params, indirect=True) +def test_marlin_fallback_no_nan_collapse(omni_runner, omni_runner_handler): + """H100 Marlin FP4 fallback: text in -> coherent text out. + + On H100 the upstream NVFP4 PWAL routes through + ``MarlinNvFp4LinearKernel.process_weights_after_loading`` -> + ``prepare_fp4_layer_for_marlin``, which casts FP8 weight_scale to + bf16 and permutes. The patched PWAL's NaN clamp runs BEFORE this + transform, so this test is the regression guard for the Marlin + code path specifically (Blackwell sm_100+ takes the FlashInfer + Cutlass/TRT-LLM/cuDNN path and is NOT covered by this test). + + Asserts the response is non-empty and does NOT match the `!!!!` + degenerate-output signature that indicates NaN propagation through + the FP4 GEMM. + """ + response = omni_runner_handler.send_omni_request({"prompts": "Why is the sky blue?", "modalities": ["text"]}) + assert response.success, "request failed" + text = (response.text_content or "").strip() + assert text, "empty response — likely server returned only EOS" + assert not _NAN_COLLAPSE_RE.search(text), ( + f"degenerate `!!!!` output detected — NaN propagated through FP4 GEMM. " + f"This means vllm_omni.patch's NVFP4 weight_scale NaN clamp did not install or " + f"did not run. Full response: {text!r}" + ) diff --git a/tests/model_executor/quantization/test_modelopt_nvfp4_nan_clamp.py b/tests/model_executor/quantization/test_modelopt_nvfp4_nan_clamp.py new file mode 100644 index 000000000..0e9f23de6 --- /dev/null +++ b/tests/model_executor/quantization/test_modelopt_nvfp4_nan_clamp.py @@ -0,0 +1,279 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for the NVFP4 W4A4 weight_scale NaN clamp installed by +``vllm_omni.patch``. + +ModelOpt 0.44's FP32->FP8 E4M3 cast of per-block weight scales can emit +literal NaN bytes (E4M3 encoding 0x7F / 0xFF) when the pre-cast scale +rounds above the FP8 max of 448 after the global-scale division. A +single NaN byte in any per-block weight_scale collapses the FlashInfer +FP4 GEMM output to NaN and the served model emits `!!!!`. The patch +overrides ``ModelOptNvFp4LinearMethod.process_weights_after_loading`` +to scan + clamp those bytes at load time. + +These tests cover the behavior of the override on a synthetic layer: + +- the NaN bytes get clamped to the FP8 E4M3 max byte (0x7E) +- a finite weight_scale is left untouched (no spurious writes) +- the install asserts itself so silent fallback is impossible +- the env-var escape hatch path is reachable + +No GPU is required — the tests construct an in-memory FP8 ``weight_scale`` +tensor with planted NaN bytes and call the patched PWAL directly. +""" + +import pytest +import torch + +# Importing vllm_omni runs vllm_omni.patch, which installs the override +# we are testing. The import must come before any direct reference to +# ``ModelOptNvFp4LinearMethod`` below. +import vllm_omni # noqa: F401 + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +FP8_NAN_BYTES = (0x7F, 0xFF) +FP8_MAX_BYTE = 0x7E # encoding of finfo(float8_e4m3fn).max + + +def _fp8_tensor_with_nan_bytes(num_blocks: int, nan_indices: list[int]) -> torch.Tensor: + """Return an FP8 E4M3 tensor where every position is the FP8 max byte + except ``nan_indices``, which carry NaN bytes (alternating 0x7F/0xFF).""" + raw = bytearray([FP8_MAX_BYTE] * num_blocks) + for i, idx in enumerate(nan_indices): + raw[idx] = FP8_NAN_BYTES[i % len(FP8_NAN_BYTES)] + # Materialize on a writable storage so masked_fill_ on the uint8 view + # can write back in-place. ``frombuffer`` produces a read-only view + # that PyTorch warns about on first use. + return torch.tensor(list(raw), dtype=torch.uint8).view(torch.float8_e4m3fn) + + +@pytest.fixture +def clamp_fn(): + """Return the top-level NaN-clamp helper exposed by vllm_omni.patch. + + Skipped when vllm doesn't expose ``ModelOptNvFp4LinearMethod`` (older + pin / unsupported build) and vllm_omni.patch consequently didn't + install the helper either. + """ + try: + from vllm_omni.patch import _clamp_nvfp4_weight_scale_nans + except ImportError: + pytest.skip("vllm_omni.patch._clamp_nvfp4_weight_scale_nans not exposed") + return _clamp_nvfp4_weight_scale_nans + + +def test_clamps_nan_bytes_in_weight_scale(clamp_fn): + weight_scale = _fp8_tensor_with_nan_bytes(num_blocks=8, nan_indices=[2, 5]) + assert torch.isnan(weight_scale).sum().item() == 2 + + layer = torch.nn.Module() + layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False) + n_clamped = clamp_fn(layer) + + assert n_clamped == 2 + assert torch.isnan(layer.weight_scale).sum().item() == 0, ( + "NaN bytes in weight_scale should have been clamped to FP8 E4M3 max" + ) + # Every byte should now be the FP8 max byte. + bytes_after = layer.weight_scale.data.view(torch.uint8).tolist() + assert bytes_after == [FP8_MAX_BYTE] * 8 + + +def test_finite_weight_scale_untouched(clamp_fn): + # All-finite weight_scale should leave bytes exactly as written. + raw = bytes([0x70, 0x71, 0x72, 0x73, 0x74, 0x75, 0x76, 0x7E]) + weight_scale = torch.frombuffer(raw, dtype=torch.uint8).clone().view(torch.float8_e4m3fn) + assert torch.isnan(weight_scale).sum().item() == 0 + before = weight_scale.data.view(torch.uint8).clone() + + layer = torch.nn.Module() + layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False) + n_clamped = clamp_fn(layer) + + assert n_clamped == 0 + after = layer.weight_scale.data.view(torch.uint8) + assert torch.equal(before, after), "clamp must not modify a finite weight_scale tensor" + + +def test_clamp_is_noop_when_weight_scale_missing(clamp_fn): + # Defensive guard: subclasses without weight_scale should not crash. + layer = torch.nn.Module() + n_clamped = clamp_fn(layer) + assert n_clamped == 0 + + +def test_clamps_non_contiguous_weight_scale(clamp_fn): + # A non-contiguous weight_scale would make `.view(torch.uint8)` raise. + # The clamp must materialize a contiguous copy, write it back to the + # Parameter, and still land the in-place clamp where the kernel reads. + # NaN bytes at even indices so they survive the stride-2 slice below + # (base[::2] keeps original indices 0, 2, 4, ...). + base = _fp8_tensor_with_nan_bytes(num_blocks=16, nan_indices=[4, 12]) + # Slice every other element -> stride 2 -> non-contiguous view. + non_contig = base[::2] + assert not non_contig.is_contiguous() + assert torch.isnan(non_contig).sum().item() >= 1 + + layer = torch.nn.Module() + layer.weight_scale = torch.nn.Parameter(non_contig, requires_grad=False) + n_clamped = clamp_fn(layer) + + assert n_clamped >= 1 + assert layer.weight_scale.is_contiguous(), "weight_scale should be contiguous after clamp" + assert torch.isnan(layer.weight_scale).sum().item() == 0, ( + "NaN bytes in a non-contiguous weight_scale should have been clamped" + ) + + +def test_install_is_self_asserting(): + # If the install assert is removed, a future regression could let the + # override silently drop. There are three legitimate end states: + # (a) upstream vLLM already has its own NaN clamp and vllm_omni.patch + # intentionally skipped installing the override (self-extinguish), + # (b) upstream is unpatched and our wrapper is installed, + # (c) the env-var escape hatch is set OR the modelopt import failed, + # so nothing is installed by design — nothing to assert. + try: + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptNvFp4LinearMethod, + ) + except ImportError: + pytest.skip("vllm.modelopt.ModelOptNvFp4LinearMethod not available") + + from vllm_omni.patch import _already_patched_upstream, _clamp_installed + + current = ModelOptNvFp4LinearMethod.process_weights_after_loading + name = getattr(current, "__qualname__", "") or getattr(current, "__name__", "") + if _already_patched_upstream: + # Self-extinguish path: presence of our marker here would mean the + # heuristic in vllm_omni.patch failed to detect upstream's clamp. + assert "_patched_nvfp4_pwal" not in name, ( + f"vllm_omni override installed despite upstream PWAL already containing a " + f"NaN clamp ({name!r}); self-extinguish heuristic in vllm_omni.patch needs review" + ) + elif _clamp_installed: + assert "patched_nvfp4_pwal" in name or "_patched_nvfp4_pwal" in name, ( + f"NVFP4 NaN-clamp PWAL override does not appear installed (saw {name!r}); " + "another plugin may have re-patched ModelOptNvFp4LinearMethod" + ) + else: + pytest.skip("NaN-clamp not installed by design (env-var escape hatch or import failure)") + + +def test_reload_does_not_nest_wrapper(): + """Reloading vllm_omni.patch must not wrap our own wrapper a second time. + + On reload the class attribute already holds our wrapper; capturing it as + the "original" and re-wrapping would run the clamp twice per load. The + patch recovers the genuine upstream method via a sentinel, so after a + reload the installed wrapper's ``_vllm_omni_wrapped_pwal`` must point at + the real upstream method (which itself carries no sentinel), not at + another wrapper. Runs in a subprocess: importlib.reload mutates global + interpreter state we don't want leaking into other tests.""" + import subprocess + import sys + import textwrap + + code = textwrap.dedent( + """ + import importlib + import vllm_omni.patch as p + try: + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptNvFp4LinearMethod, + ) + except ImportError: + print("SKIP: modelopt not available") + raise SystemExit(0) + if not getattr(p, "_clamp_installed", False): + print("SKIP: clamp not installed (upstream patched or escape hatch)") + raise SystemExit(0) + importlib.reload(p) + fn = ModelOptNvFp4LinearMethod.process_weights_after_loading + wrapped = getattr(fn, "_vllm_omni_wrapped_pwal", None) + # The recovered upstream must NOT itself be one of our wrappers. + nested = hasattr(wrapped, "_vllm_omni_wrapped_pwal") + print("NESTED" if nested else "FLAT") + """ + ) + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=120) + if result.stdout.strip().startswith("SKIP:"): + pytest.skip(result.stdout.strip()) + assert "FLAT" in result.stdout, ( + f"reload nested the NVFP4 clamp wrapper — clamp would run twice per load.\n" + f"stdout: {result.stdout}\nstderr: {result.stderr}" + ) + + +@pytest.mark.parametrize( + ("env_value", "acceptable"), + [ + # Truthy values activate the escape hatch: the import short-circuits + # before the upstream-detection runs, so nothing is installed. + ("1", frozenset({"NOT_INSTALLED"})), + ("true", frozenset({"NOT_INSTALLED"})), + # Falsy values mean "do NOT use the escape hatch" — the repo-wide + # bool-env idiom in vllm_omni.patch guards against a naive + # `if os.environ.get(...)` wrongly skipping on "0"/"false". The clamp + # must therefore be in force, which is satisfied EITHER by our wrapper + # being installed OR by upstream vLLM already carrying its own clamp + # (self-extinguish, `_already_patched_upstream`), in which case skipping + # our wrapper is the correct behavior. Accept both so this test keeps + # passing once the upstream fix lands. + ("0", frozenset({"INSTALLED", "UPSTREAM_PATCHED"})), + ("false", frozenset({"INSTALLED", "UPSTREAM_PATCHED"})), + ], +) +def test_env_var_escape_hatch_disables_install(env_value, acceptable): + """Spawn a subprocess with VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP set and verify + the override install state matches the bool-env convention. The patch + module runs at ``vllm_omni`` import time, so we can't toggle it via + monkeypatch in the same process. + + The subprocess reports the authoritative state from vllm_omni.patch's + module flags (the same source of truth as ``test_install_is_self_asserting``) + rather than introspecting the method qualname — both because that is the + real contract and because ``"INSTALLED"`` is a substring of + ``"NOT_INSTALLED"`` and naive substring matching would be vacuous.""" + import os + import subprocess + import sys + import textwrap + + code = textwrap.dedent( + """ + import vllm_omni # noqa: F401 + import vllm_omni.patch as p + try: + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptNvFp4LinearMethod, # noqa: F401 + ) + except ImportError: + print("SKIP: modelopt not available") + raise SystemExit(0) + if getattr(p, "_already_patched_upstream", False): + # Upstream vLLM already clamps; vllm_omni.patch deliberately + # skipped installing our wrapper (self-extinguish). + print("UPSTREAM_PATCHED") + elif getattr(p, "_clamp_installed", False): + print("INSTALLED") + else: + print("NOT_INSTALLED") + """ + ) + env = os.environ.copy() + env["VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP"] = env_value + # 120s rather than 60s: cold CI containers can take 30-50s just to + # import vllm_omni → vllm → torch before the subprocess does any work. + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, timeout=120, env=env) + if result.stdout.strip().startswith("SKIP:"): + pytest.skip("modelopt not available in subprocess") + # Exact token match on the final stdout line — not a substring test, since + # "INSTALLED" is a substring of "NOT_INSTALLED". + state = result.stdout.strip().splitlines()[-1] if result.stdout.strip() else "" + assert state in acceptable, ( + f"VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP={env_value!r} expected one of {set(acceptable)!r} " + f"but got {state!r}:\nstdout: {result.stdout}\nstderr: {result.stderr}" + ) diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index 59d697c16..8bccb324f 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -1,7 +1,9 @@ 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.inputs import TokensPrompt as _OriginalTokensPrompt @@ -21,6 +23,8 @@ 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") + # ============================================================================= # Patch ModelConfig.is_mm_prefix_lm to support omni-specific models # ============================================================================= @@ -71,6 +75,185 @@ assert _installed is _patched_cp, ( "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: ModelOptNvFp4LinearMethod (W4A4 NVFP4 Linear) only. NvFp4FusedMoE / +# NvFp4W4A16 / CompressedTensors / Quark NVFP4 paths are not covered. +# +# SELF-EXTINGUISH: `_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 is expected to take when it is +# filed (planned as a follow-up PR after this one merges). Once vllm-omni's +# vllm pin moves to a release with that upstream fix, the override is +# skipped at import and this block can be deleted. NOTE: the heuristic only +# matches "vLLM PWAL clamps NaN in-place with `masked_fill_`"; if the +# 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 — +# but the heuristic should be revisited when the upstream PR is filed. +# +# 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 PWAL contains its own NaN clamp. +# `_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 being patched, so the +# right check for "we own NaN-clamp behavior" is `_clamp_installed`. +_already_patched_upstream = False +_clamp_installed = False + +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. + # 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") + from vllm.model_executor.layers.quantization.modelopt import ( + ModelOptNvFp4LinearMethod as _OriginalModelOptNvFp4LinearMethod, + ) +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, + ) +else: + _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 # =============================================================================