[Bugfix] Clamp NVFP4 MoE activation scales to prevent NaN from dead experts

[pavelzak]

Summary

Some NVFP4 checkpoints (e.g. txn545/Qwen3.5-122B-A10B-NVFP4) have near-zero or denormal float32 activation scales for dead MoE experts. In make_nvfp4_moe_quant_config(), these are inverted directly (1.0 / a13_scale, 1.0 / a2_scale), producing Inf, which the FP4 quantization kernel propagates as NaN — corrupting the entire output batch.

Fix: clamp both activation global scales from below at torch.finfo(torch.float32).tiny (~1.18e-38) before inversion. Dead experts contribute zero-weighted outputs regardless of their scale value, so this clamp is numerically safe.

• This is not a duplicate of feat: Add SM121/GB10 (DGX Spark) Blackwell-class GPU support #31740 (which covers SM12.1 platform support); that PR does not touch oracle/nvfp4.py.

Test plan

• Loaded txn545/Qwen3.5-122B-A10B-NVFP4 on DGX Spark (SM12.1, CUDA 13.1) with vLLM v0.20.2 + this patch
• Ran multi-turn chat completions — no NaN/Inf in outputs
• Without the patch, outputs contain NaN tokens when dead experts are routed

AI assistance was used to help identify and describe this fix.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @pavelzak.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[gemini-code-assist]

Code Review

This pull request introduces extensive updates to vLLM, including automated Docker image publishing, support for input clamping in SwiGLU activations, and a fused hc_head kernel implemented with TileLang. Significant optimizations are added for DeepSeek-V4, such as multi-stream GEMM overlap and expert-dtype-aware MoE dispatch. Additionally, the V1 engine's KV cache management is refined to better handle block recycling for sliding window and chunked-local attention. Review feedback correctly identified critical bugs in the DeepSeek-V4 attention layer where ReplicatedLinear outputs were not unpacked from their return tuples before being reshaped, which would cause attribute errors at runtime.

I am having trouble creating individual review comments. Click here to see my feedback.

vllm/model_executor/layers/deepseek_v4_attention.py (427)

The ReplicatedLinear layer returns a tuple of (output, bias). Attempting to call .view() directly on the return value will raise an AttributeError because tuple does not have a view method. You must unpack the output first.

                q, _ = self.wq_b(qr)
                q = q.view(-1, self.n_local_heads, self.head_dim)

vllm/model_executor/layers/deepseek_v4_attention.py (453)

Similar to the issue above, self.wq_b(qr) returns a tuple (output, bias). You need to unpack the output before calling .view().

                q, _ = self.wq_b(qr)
                q = q.view(-1, self.n_local_heads, self.head_dim)

vllm/model_executor/layers/deepseek_v4_attention.py (466)

The ReplicatedLinear layer returns a tuple. Unpack the output before applying the view transformation.

            q, _ = self.wq_b(qr)
            q = q.view(-1, self.n_local_heads, self.head_dim)

[henry2-9]

Confirmed the bug and verified this fix end-to-end on DGX Spark (GB10, sm_121, CUDA 13.0). Wanted to share a regression test since the PR doesn't currently include one — happy to
open a follow-up PR adding it to tests/kernels/moe/test_nvfp4_moe.py after this merges.

Repro on stock vllm==0.20.2

import pytest
import torch
from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
    NvFp4MoeBackend,
    make_nvfp4_moe_quant_config,
)

# All backends that hit the buggy `1.0 / a13_scale` inversion
# (MARLIN and EMULATION take separate branches).
_BUGGY_PATH_BACKENDS = [
    NvFp4MoeBackend.VLLM_CUTLASS,
    NvFp4MoeBackend.FLASHINFER_CUTLASS,
    NvFp4MoeBackend.FLASHINFER_TRTLLM,
    NvFp4MoeBackend.FLASHINFER_CUTEDSL,
    NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED,
]


@pytest.mark.parametrize("backend", _BUGGY_PATH_BACKENDS)
def test_make_nvfp4_moe_quant_config_clamps_dead_expert_scales(backend):
    """Regression test for #42601 — dead-expert scales must not produce inf/NaN."""
    E = 8 
    # Two dead experts in a13 (indices 0, 3), two dead in a2 (indices 2, 7).
    a13_scale = torch.tensor(
        [0.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0], dtype=torch.float32
    )
    a2_scale = torch.tensor(
        [1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0, 0.0], dtype=torch.float32
    )
    w13_scale_2 = torch.ones((E,), dtype=torch.float32)
    w2_scale_2 = torch.ones((E,), dtype=torch.float32)
    w13_scale = torch.ones((E, 1, 1), dtype=torch.uint8)
    w2_scale = torch.ones((E, 1, 1), dtype=torch.uint8)

    config = make_nvfp4_moe_quant_config(
        backend=backend,
        w13_scale=w13_scale,
        w2_scale=w2_scale,
        w13_scale_2=w13_scale_2,
        w2_scale_2=w2_scale_2,
        a13_scale=a13_scale,  
        a2_scale=a2_scale,
    )

    assert torch.isfinite(config.a1_gscale).all(), (
        f"[{backend.name}] a1_gscale must be finite for dead-expert scale=0.0; "
        f"got {config.a1_gscale}"
    )
    assert torch.isfinite(config.a2_gscale).all()


def test_make_nvfp4_moe_quant_config_healthy_scales_unchanged():
    """Sanity check: clamp is a no-op when all scales are healthy."""
    E = 4
    a13_scale = torch.tensor([0.5, 1.0, 2.0, 4.0], dtype=torch.float32)
    a2_scale = torch.tensor([1.0, 2.0, 0.5, 4.0], dtype=torch.float32)
    config = make_nvfp4_moe_quant_config(
        backend=NvFp4MoeBackend.VLLM_CUTLASS,
        w13_scale=torch.ones((E, 1, 1), dtype=torch.uint8),
        w2_scale=torch.ones((E, 1, 1), dtype=torch.uint8),
        w13_scale_2=torch.ones((E,), dtype=torch.float32),
        w2_scale_2=torch.ones((E,), dtype=torch.float32),
        a13_scale=a13_scale,
        a2_scale=a2_scale,
    )
    torch.testing.assert_close(config.a1_gscale, 1.0 / a13_scale)
    torch.testing.assert_close(config.a2_gscale, 1.0 / a2_scale)

Results

Without this PR (stock vllm==0.20.2): 5 of 6 fail; sanity-check passes.

FAILED test_make_nvfp4_moe_quant_config_clamps_dead_expert_scales[NvFp4MoeBackend.VLLM_CUTLASS]
   ... got tensor([inf, 1., 1., inf, 1., 1., 1., 1.])
FAILED ...FLASHINFER_CUTLASS
FAILED ...FLASHINFER_TRTLLM
FAILED ...FLASHINFER_CUTEDSL
FAILED ...FLASHINFER_CUTEDSL_BATCHED
PASSED  test_make_nvfp4_moe_quant_config_healthy_scales_unchanged
== 5 failed, 1 passed in 2.84s ==

With this PR applied (clamp at torch.finfo(float32).tiny): all 6 green.

== 6 passed in 2.82s ==

Tested on:

• NVIDIA GB10 (DGX Spark, sm_121), CUDA 13.0, vLLM 0.20.2 + this PR cherry-picked
• The healthy_scales_unchanged sanity case confirms the clamp threshold (~1.18e-38) is below any realistic activation scale, so it's a true no-op for normal checkpoints.

Glad to open the follow-up PR after this merges — looks good to me, +1!

cc @pavelzak

---

[rjmayott]

Confirmed both bug and fix on SM120.

[pavelzak]

Could a maintainer add the ready label to trigger CI? The fix has been validated on DGX Spark (@henry2-9) and SM120 (@rjmayott). Happy to address any review feedback.

[pavelzak]

Friendly ping — this has been open since May and still shows mergeable against main. The fix has since been battle-tested further: it's been running in production on a 2× DGX Spark (GB10) TP=2 cluster serving DeepSeek-V4-Flash-0731 (NVFP4 experts) as part of the stack described in #52499–#52503; without the clamp, dead experts produce NaN activations on this checkpoint. Happy to rebase if maintainers prefer a fresh CI run.

[Jie-Fang]

• clamp(min=tiny) does not fix NaN: if the input is NaN, the output remains NaN.

• +Inf is inverted to 0, while negative scales are silently clamped to tiny, masking these invalid values.

• The current ModelOpt path allocates scales using torch.empty. A missing value is not guaranteed to be zero; it may contain an arbitrary finite garbage value, which this PR cannot detect at all.

Refer to #45320 / #54444, they have a better way to deal with such situation.

And input_scale will be used in:

a_gscale = 1.0 / input_scale
weight_scale_2 *= input_scale

if the input_scale is zero, the original path will make a_gscale inf and weight_scale_2 0, the result is inf * 0 = NaN which can be observed by the numeric issue. But this PR makes the result become zero, how can we make sure it's a good number?

[pavelzak]

Thanks for the review. I agree that making the reciprocal finite does not establish output correctness, and my numerical-safety claim was too strong. I no longer use the affected model and cannot commit to the additional validation needed, so I’m closing this workaround. The load-time validation approach in #45320 / #54444 is a better direction for missing or invalid scales.