PR #39391: fix: clamp NaN/Inf in topk_softmax to prevent duplicate expert IDs

• Repository: vllm-project/vllm
• Author: jhaotingc
• State: closed | merged: True
• Link: https://github.com/vllm-project/vllm/pull/39391

Description (problem / solution / changelog)

Purpose

Fix https://github.com/vllm-project/vllm/issues/39244

Fix CUDA illegal memory access crash when serving MoE models (e.g., Qwen3.5-397B-A17B-FP8) with FlashInfer CUTLASS MoE and CUDA graphs at high concurrency on H200.

CUDA graph replay pads the batch to the nearest capture size. Padded tokens have degenerate hidden states that produce NaN gating logits. The topkGating kernel's softmax outputs all-NaN, and the argmax loop picks expert 0 for every top-k slot (IEEE 754: NaN > NaN is false), producing duplicate expert IDs [0,0,0,0,0,0,0,0]. These duplicates trigger an uninitialized-memory bug in FlashInfer's three-step MoE sort, causing finalizeMoeRoutingKernel to dereference wild pointers.

The fix clamps NaN/Inf values to 0 after softmax/sigmoid scoring in topkGating, before the argmax selection loop. With all-zero scores, argmax picks unique experts [0,1,2,...,k-1] via index tie-breaking. Zero performance overhead.

Test Plan

• Kernel unit test: verify topk_softmax produces unique expert IDs for NaN/Inf/normal gating inputs

• Kernel microbenchmark: compare eager + CUDA graph replay latency for normal vs NaN inputs (batch 1-512, 128/256 experts)

• End-to-end: serve Qwen3.5-397B-A17B-FP8 (TP=4, EP=4, CUDA graphs, VLLM_USE_FLASHINFER_MOE_FP8=1) with 8 concurrent requests

• Full sweep: sglang benchmark conc 1-512, ISL=1600, OSL=600, REPEAT=5 on 4x H200

• tests/kernels/moe/test_fused_topk.py::test_fused_topk_nan_inf_clamp

Test Result

Kernel correctness (H200):

Kernel perf (H200, CUDA graph replay, median of 1000 runs):

All within noise. Zero measurable overhead.

End-to-end (4x H200, Qwen3.5-397B-A17B-FP8):

unit test

Summary

vLLM crashes with CUDA error: an illegal memory access was encountered when serving Qwen3.5-397B-A17B-FP8 with VLLM_USE_FLASHINFER_MOE_FP8=1 and CUDA graphs enabled. The crash occurs at high concurrency (8+ requests) when the MoE batch size exceeds 256 tokens.

Root Cause

CUDA graph replay pads the batch to the nearest capture size (e.g., 300 real tokens padded to 512). Padded tokens have stale/degenerate hidden states that produce NaN gating logits in the MoE router. The topk_softmax CUDA kernel then produces duplicate expert IDs for NaN inputs (e.g., [0,0,0,0,0,0,0,0] for every padded token), because IEEE 754 NaN > NaN is always false, so the argmax never updates from expert 0, and the -10000 zeroing of the winner also fails (-10000 > NaN is false).

These duplicate expert IDs trigger a latent bug in FlashInfer's blockExpertPrefixSumKernel (three-step MoE sort path, used when num_tokens > 256): it uses break after the first expert match, so duplicate expert slots leave unpermuted_row_to_permuted_row entries uninitialized. finalizeMoeRoutingKernel then reads garbage values as row indices, causing wild pointer dereferences.

Chain of events

CUDA graph replay with padded tokens
  -> stale hidden states -> NaN gating logits
    -> topk_softmax produces [0,0,0,0,0,0,0,0] for padded tokens
      -> duplicate expert IDs enter cutlass_fused_moe (num_tokens > 256)
        -> blockExpertPrefixSumKernel skips duplicate slots (break)
          -> unpermuted_row_to_permuted_row has uninitialized entries
            -> finalizeMoeRoutingKernel reads garbage -> OOB -> CRASH

Why it only happens with CUDA graphs

In eager mode, there are no padded tokens -- the batch contains only real tokens with valid hidden states, the router produces unique expert IDs, and the three-step sort works correctly. The crash requires:

1. Batch size > 256 (three-step sort path)
2. Duplicate expert IDs (from NaN gating on padded tokens)

Both conditions only occur together during CUDA graph replay at high concurrency.

Fix

Clamp NaN/Inf values to 0 in topk_softmax after softmax/sigmoid scoring, before the argmax selection loop:

// csrc/moe/topk_softmax_kernels.cu, after line 443
#pragma unroll
for (int ii = 0; ii < VPT; ++ii) {
    if (isnan(row_chunk[ii]) || isinf(row_chunk[ii])) {
        row_chunk[ii] = 0.f;
    }
}

With all-zero scores, the argmax uses index tie-breaking to pick unique experts [0,1,2,...,k-1], preventing duplicates. Normal (non-NaN) inputs are unaffected -- the clamp is a no-op.

Why this is the right fix location

The topk_softmax kernel (csrc/moe/topk_softmax_kernels.cu:266) is where the NaN propagates into duplicate expert IDs. Fixing it here:

• Prevents the bad input from reaching ANY downstream MoE kernel (FlashInfer, Triton, etc.)
• Zero performance overhead (see benchmarks below)
• Handles all NaN sources (CUDA graph padding, numerical overflow, any future degenerate input)

Performance Impact

Benchmarked on H200, production MoE configs (128/256 experts, top_k=8). The fix adds isnan/isinf checks (single PTX predicate instructions) per element. The kernel is memory-bandwidth bound, so the extra comparisons are invisible:

Eager mode (us, median of 1000 runs)

CUDA graph replay mode (us, median of 1000 runs)

All differences are within noise (<2%). Zero measurable overhead.

Verification

Standalone (topk kernel)

Before fix:

all_nan: dup_tokens=512/512  topk_ids=[0,0,0,0,0,0,0,0]
all_inf: dup_tokens=512/512  topk_ids=[0,0,0,0,0,0,0,0]

After fix:

all_nan: dup_tokens=0/512  topk_ids=[0,1,2,3,4,5,6,7]
all_inf: dup_tokens=0/512  topk_ids=[0,1,2,3,4,5,6,7]

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

Changed files

• csrc/moe/topk_softmax_kernels.cu (modified, +19/-2)
• tests/kernels/moe/test_fused_topk.py (modified, +67/-0)

PR #40553: test: add nan/inf clamp regression test for fused_topk_bias

• Repository: vllm-project/vllm
• Author: jhaotingc
• State: closed | merged: True
• Link: https://github.com/vllm-project/vllm/pull/40553

Description (problem / solution / changelog)

Add test_fused_topk_bias_nan_inf_clamp to cover the same NaN/Inf scenario as test_fused_topk_nan_inf_clamp but for the fused_topk_bias entry point. On CUDA, fused_topk_bias routes through the same topk_softmax/topk_sigmoid kernels fixed in #39391 (lines 131/154 of topk_softmax_kernels.cu), so the clamp already applies.

Closes #40457

Purpose

PR #39391 fixed NaN/Inf gating logits producing duplicate expert IDs in fused_topk by clamping scores to 0 in topk_softmax_kernels.cu (lines 131/154 for the warp-level path, line 457 for the fallback path). However, the regression test added in that PR only covered fused_topk. This PR adds the same nan/inf regression test for fused_topk_bias, which is used by DeepSeek-style models with e_score_correction_bias.

On CUDA, fused_topk_bias routes through the same topk_softmax / topk_sigmoid C++ kernels that were patched in #39391, so the clamp already applies. This PR confirms that coverage with an explicit test.

Test Plan

• Added test_fused_topk_bias_nan_inf_clamp to tests/kernels/moe/test_fused_topk.py, parametrized over:
  • dtype: bfloat16, float16, float32
  • scoring_func: softmax, sigmoid
  • bad_value: NaN, Inf
  • num_experts: 6, 8, 16
  • topk: 3, 4
  • Total: 144 test cases
• Verified the test is automatically collected by the existing Kernels MoE Test CI step (no .buildkite/ changes needed):

  # Simulate exactly what buildkite runs:
  pytest tests/kernels/moe --collect-only -q | grep nan_inf
  # Shows both test_fused_topk_nan_inf_clamp (added in #39391) and
  # test_fused_topk_bias_nan_inf_clamp (this PR) are collected

• Ran the full test_fused_topk.py suite (720 tests) to check for regressions.

Test Result

nan/inf regression tests (144 cases)

Full test suite (720 cases, 2×H200, CUDA 13.0)

720 passed, 16 warnings in 104.91s

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Changed files

• tests/kernels/moe/test_fused_topk.py (modified, +69/-0)