[Bug] DeepSeek-V4 on SM12x: NaN MQA logits drive top_k_per_row_prefill to emit uninitialized smem as indices -> illegal memory access

[linjiapro]

Environment

• vLLM 0.26.0 (release wheel), flashinfer-python 0.6.14, torch 2.11.0+cu130, CUDA 13.0, driver 580.126.09
• 4× NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition (SM120), TP=4
• Model: deepseek-ai/DeepSeek-V4-Flash, --kv-cache-dtype fp8 --block-size 256 --tokenizer-mode deepseek_v4 --reasoning-parser deepseek_v4 --trust-remote-code, OpenAI-compatible server

Symptom

DeepSeek-V4-Flash on SM120 (enabled by #43477) crashes with CUDA error: an illegal memory access was encountered on all TP ranks → EngineDeadError, on the first prefill batch containing a request longer than ~2048 tokens (e.g. 16 concurrent structured-output requests with 0.5–8K-token prompts; also a single ~5K-token prompt). Short prompts serve perfectly. We reproduced the crash 9 times across vLLM 0.25.1/0.26.0 and flashinfer 0.6.13/0.6.14/0.6.15.

Under CUDA_LAUNCH_BLOCKING=1, the fault is raised in the Triton kernel _compute_global_topk_indices_and_lens_kernel (vllm/models/deepseek_v4/common/ops/cache_utils.py), reached from DeepseekV4FlashInferSM120Attention._forward_prefill (compress-ratio-4 branch).

Root-cause chain (instrumented at the tensor level)

1. fp8_fp4_mqa_logits produces NaN-dominated logits on SM120.
Auditing the valid [ks, ke) region of the prefill MQA logits per C4A layer on a single 4,972-token prefill: the first C4A layer is 83% NaN (2,560,388 of 3,088,855 entries; the finite remainder looks plausible, range ≈ −2…4), and every subsequent layer is 100% NaN (cascade through the model). This is presumably the SM12x scale-handling issue (cf. the e8m0→fp32 upcast in #41834).

2. top_k_per_row_prefill mishandles NaN input on its histogram path and emits uninitialized shared memory as indices.
Auditing topk_indices_buffer immediately after ops.top_k_per_row_prefill returns: rows whose candidate count rowLen = ke−ks ≤ topk (the short-row path) are clean (-1-padded as documented); rows with rowLen > topk (histogram path) contain garbage starting at exactly the first such row, e.g. rows 2051…4868 of a 4,869-token request at compress ratio 4 (2051/4 ≈ 513 > 512 = index_topk). Garbage row counts differ across TP ranks on identical input (1923/2151/2120/1987) — nondeterministic, consistent with uninitialized dynamic shared memory (smemOutput is never pre-initialized and the final copy loop writes all topK slots).

This reproduces in isolation, without serving (25 lines, below): with clean random logits the kernel is correct on SM120 (0 garbage rows, and tests/kernels/test_top_k_per_row.py passes 12/12 here); with 83%-NaN logits, ~40% of histogram-path rows contain garbage. Two distinct failure signatures:

• indices equal to 2143289344 = 0x7FC00000 — a float32 NaN bit pattern reinterpreted as an int32 index;
• output slots that keep their pre-launch canary value — i.e. never written at all (in serving these become whatever is in shared memory).

Since the repro is data-driven (NaN input) rather than SM120-specific, the kernel's NaN mishandling likely affects all architectures; SM120 is where NaN logits actually occur in practice.

import torch, vllm  # importing vllm registers the _C ops
torch.manual_seed(0); torch.set_default_device("cuda:0")

rows, cols, top_k = 4869, 1243, 512
ks = torch.zeros(rows, dtype=torch.int32)
ke = torch.clamp(torch.arange(rows, dtype=torch.int32) // 4 + 1, max=cols)

def probe(tag, logits):
    out = torch.full((rows, top_k), -7, dtype=torch.int32)  # -7 = canary
    torch.ops._C.top_k_per_row_prefill(logits, ks, ke, out, rows,
                                       logits.stride(0), logits.stride(1), top_k)
    torch.cuda.synchronize()
    garb = (out >= 1_000_000) | ((out < -1) & (out != -7))
    print(f"{tag}: garbage_rows={int(garb.any(dim=1).sum())}/{rows} "
          f"unwritten_rows={int((out == -7).any(dim=1).sum())} max={int(out.max())}")

logits = torch.randn(rows, cols)
probe("clean ", logits.clone())                                  # 0 garbage
nanl = logits.clone(); nanl[torch.rand(rows, cols) < 0.83] = float("nan")
probe("83%NaN", nanl)   # ~1800 garbage rows, ~500 unwritten, max=2143289344

3. compute_global_topk_indices_and_lens dereferences the garbage.
Its validity mask is only local_idx >= 0, so INT32_MAX-class indices pass and it gathers block_table[req, idx // block_size] far out of bounds → the illegal memory access that kills the engine.

Behavioral confirmation

Needle-in-haystack on the same server: a prompt below the sparse threshold retrieves the needle perfectly ("The secret code word is BLUEBERRY."); the identical needle in a ~4.5K-token prompt yields output completely unrelated to the prompt (the model hallucinates arbitrary code/text — its selected KV set is effectively random). Short prompts mask defect (1) entirely because when rowLen ≤ topk the top-k selects all candidates and scores never matter.

Suggested fixes (layered)

1. Kernel hardening (csrc/libtorch_stable/sampler.cu): the histogram path should never emit unwritten smem or NaN bit patterns as indices — pad under-filled output with -1 like the short-row path, and exclude NaN scores from selection. Data-driven robustness fix, valuable on every arch. (We can't build csrc locally to contribute this ourselves.)
2. SM12x routing: until (1) lands, route SM12x prefill top-k through a torch.topk implementation of the kernel's documented contract — the decode path already excludes SM12x from cooperative_topk. PR incoming; verified on the affected hardware that the previously 9×-reproducible crash is eliminated.
3. The NaN logits themselves (the primary SM12x defect): DeepGEMM fp8_fp4_mqa_logits on SM120. [New Model][Nvidia] Add SM12x support for DeepSeek V4 Flash with essential fixes #41834 appears to carry the relevant fixes (e8m0→fp32 upcast, SM120 DeepGEMM fallbacks) but is unmerged; even with (1)+(2), DSv4 output quality on SM120 remains broken for >2048-token prompts until this is addressed — with NaN scores, any correct top-k selects an arbitrary candidate set.

Happy to run further diagnostics on this hardware — SM120 workstations seem rare in CI.

[linjiapro]

Additional evidence isolating stage (1) — the NaN logits — to a race condition in the DeepGEMM MQA logits kernel on SM120, established by testing DeepGEMM nv_dev at head (f8e8fb5, deep_gemm 2.6.1) standalone on the same 4× RTX PRO 6000 Blackwell Workstation box (built with g++-10, installed so vLLM prefers it over the vendored copy):

DeepGEMM's own test suite fails on SM120 for this kernel family, in two independent ways:

1. tests/test_attention.py::test_mqa_logits, fp8 format (the config DSv4-Flash serving actually uses — vLLM logs "Using FP8 indexer cache for Lightning Indexer", index_n_heads=64, head_dim=128): the test's self-consistency check fails — the same kernel on identical input across 20 repeat launches returns bitwise-different results:

   AssertionError: bitwise mismatch (mqa logits self-consistency): num_bytes=1048576,
   num_mismatch=56571, ... x_val=0.0, y_val=1.046875
   

   i.e. the kernel is nondeterministic on SM120 — a synchronization bug. This lineage has needed per-arch sync fixes before (SM90: Fix SM90 MQA logits deepseek-ai/DeepGEMM#229, SM100: Fix a sync issue in SM100 MQA logits deepseek-ai/DeepGEMM#285 "Fix a sync issue in SM100 MQA logits", plus 96fbced "Fix SM100 FP16 MQA synchronization"); SM120 appears to be missing its equivalent.

2. Same test, mxfp4 format (H=64, D=128 — the FP4-indexer config): accuracy failure, Diff: 0.6187 vs 0.02 tolerance, at full speed (~776 TFLOPS) — wrong numbers, not a crash.

A racy accumulator/flag read also explains the serving-side observations better than a pure scale-format bug: NaN fractions varying per TP rank on identical input, and 83% → 100% NaN progression across layers.

Practical status for SM12x DSv4 after #49897: engine no longer crashes, but output for >2048-token prompts remains incoherent until the DeepGEMM kernel gets an SM120 sync fix (or vLLM's vendored copy does). The 5-minute repro on affected hardware is just:

python -c "
import sys; sys.path.insert(0, '<DeepGEMM>/tests')
import test_attention as ta
orig = ta.sample_mqa_cases
ta.sample_mqa_cases = lambda tag, cases: [c for c in orig(tag, cases) if c[0] == 'fp8']
ta.test_mqa_logits()"

[linjiapro]

The DeepGEMM-side defect is now filed against the branch that owns SM120 support (nv_dev — main has no SM12x sources at all): deepseek-ai/DeepGEMM#392, with the self-consistency repro. Once a sync fix lands there, the remaining work on the vLLM side is picking it up in the vendored copy (or preferring a fixed standalone install).

[linjiapro]

Status update after a much deeper isolation pass on the same hardware — the causal chain needs one correction and gains two new findings:

Correction: the DeepGEMM MQA logits kernel is not the NaN source. Instrumenting the inputs to fp8_fp4_mqa_logits at the first C4A layer during a failing long prefill: q quant is clean, gathered kv is clean, scales are clean — but the indexer weights are ~71% NaN, and one step further up, hidden_states entering the layer is already 65–78% NaN (per-TP-rank varying on identical input, i.e. nondeterministic). The NaN is born in layers 0–1 (the SWA-only, compress-ratio-0 layers) during prefills longer than ~2K tokens, before any sparse-indexer machinery runs. Probing the DeepGEMM kernel standalone (nv_dev @ f8e8fb5) confirms its valid-region output is deterministic and correct; its only defect is nondeterministic out-of-window stores (tracked at deepseek-ai/DeepGEMM#392) which are invisible to consumers that respect the row windows.

Everything in the original report about top_k_per_row_prefill stands unchanged: given NaN-dominated logits it emits NaN bit patterns and uninitialized smem as indices (standalone repro in the issue body), and #49897 remains the right guard for that layer.

New finding 1: the vendored vllm/third_party/deep_gemm (a ~2.5.0 nv_dev snapshot) contains a paged MQA decode bug fixed upstream on nv_dev: the TMA producer prefetches q_idx + last_advance, which "may point at the wrong Q tile" when zero-context atoms are skipped (mixed prefill+decode batches — exactly vLLM's serving shape). nv_dev's fix in sm120_fp8_paged_mqa_logits.cuh ("Use the scheduler's actual next task") looks like a clean cherry-pick into the vendored copy. (Swapping the whole standalone 2.6.1 in is not viable: its grouped MoE GEMM launch-fails under 0.26.0, and the paged path depends on the vendored-vintage cache layout written by vLLM's own indexer_k_quant_and_cache.)

New finding 2 / open question: the remaining root cause on SM12x is whatever corrupts hidden_states in the first two (SWA-only) decoder layers on long prefill — per-rank nondeterministic, length-gated (~2K+). The suspects that #41834 patches in this area (fused_deepseek_v4_qnorm_rope_kv_insert_kernel.cu, cache_kernels.cu, the SWA prefill path) seem consistent with this. Happy to run further bisection on this hardware.

[linjiapro]

Major methodological finding plus five eliminated mechanisms, from a bisection session on the affected hardware:

The corruption is stochastic per server launch, and sticky within a launch. One launch out of ~10 served long prompts perfectly (needle retrieval PASS at every length up to 5.2K tokens, 10/10 sequential probes); every other launch fails every long prompt, including with identical request parameters (temp 0, thinking off). Within a launch the state never changes. Any single-launch A/B experiment on this bug is therefore unreliable — including some of our own earlier inferences.

Mechanisms eliminated (each tested by toggling one variable and re-running the long-prompt needle battery, multiple launches where feasible):

• Sampling parameters / reasoning mode (fail and pass both reproduce across all combinations on a given launch).
• DeepGEMM PDL (set_pdl(False)): still fails.
• VLLM_USE_BREAKABLE_CUDAGRAPH=0: still fails, 3/3 launches.
• Full enforce_eager (no CUDA graphs at all): still fails — capture machinery exonerated.
• MHC hyper-connection prenorm via DeepGEMM tf32_hc_prenorm_gemm (forced the TileLang fallback): still fails.

Combined with the earlier tensor-level audits (NaN present in hidden_states at the first C4A layer's entry; per-TP-rank varying counts on identical input), the surviving suspect set for the layer-0/1 long-prefill corruption is narrow:

1. fused_deepseek_v4_qnorm_rope_kv_insert (csrc; patched by [New Model][Nvidia] Add SM12x support for DeepSeek V4 Flash with essential fixes #41834),
2. the SM120 SWA prefill attention path (flashinfer_sparse SWA branch),
3. the FP8 block-scaled linear path (DeepGemmFp8BlockScaledMMKernel; UE8M0 checkpoint pins linear to DeepGEMM) — possibly with launch-time allocation-address / alignment dependence, which would explain the per-launch stickiness.

The launch-stochastic character suggests an address/alignment-conditional defect in one of those kernels (allocations are fixed at startup, which is exactly the per-launch state that differs). Happy to continue bisection — the next step is per-layer NaN instrumentation on a known-poisoned launch to separate (1)/(2)/(3).

[Smallfu666]

I can build csrc and have H100/H200 hardware to validate on, so I'd like to take the Layer-1 kernel hardening slice of this: NaN exclusion in the histogram selection path plus smemOutput initialization/padding in csrc/libtorch_stable/sampler.cu, kept separate from #49897 and the DeepGEMM-side fix. Plan: reproduce with the canary repro above (failing baseline first), then a small kernel patch and a NaN-input case in tests/kernels/test_top_k_per_row.py. Will post a PR if the fix validates.