[Quantization][Qwen3-Omni] Enable NVFP4 W4A4 serving on Blackwell (#4025)

Signed-off-by: Yihong Jin <yihongjin404@gmail.com>

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

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}"
+    )

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}"
+    )

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
 # =============================================================================