[Fix] SM120 sparse MLA: zero-initialise the 64-token page-split scratch so masked candidates never gather NaN from slot 0

[avifenesh]

Motivation

On SM120 (RTX PRO 6000 Blackwell) with the FlashInfer sparse-MLA backend (SGLANG_SM120_FLASHMLA_BACKEND=flashinfer, the default) and the pinned FlashInfer 0.6.18, DeepSeek-V4 family models return wrong attention for prefill calls with more than 64 query rows. Observed on DeepSeek-V4.1-Flash, TP4, 4x RTX PRO 6000:

• exact-fit oracle over 20 cold prompts: 60 and 64 prompt tokens -> 8/8 correct; 65 / 66 / 80 tokens -> 0/12 (garbage; prompt_tokens confirmed via /tokenize and usage);
• planted-code recall sweep over 23 cold prompts of 36..32k tokens: the code is retrieved 4/23 on one box and 5/23 on an independent one, only for the prompts of 64 tokens or fewer; tool-call arguments come back as {}, vision descriptions are garbage;
• the same sweep is 23/23 with the Triton implementation (SGLANG_SM120_FLASHMLA_BACKEND=triton), which is how the defect was first contained.

The boundary is SM120_DECODE_MAX_TOKENS = 64 in python/sglang/kernels/ops/attention/flash_mla_sm120.py: at most 64 rows go to the FlashInfer decode (split-K) kernel, more rows go to the FlashInfer prefill kernel via _flash_mla_sm120_prefill.

Root cause

An interaction between the FlashInfer kernels and the way sglang feeds them the KV pages:

1. FlashInfer <= 0.6.18 gathers slot 0 for masked candidates. The SM120 sparse-MLA prefill and decode kernels clamp every padding index (-1) to slot 0 and gather that slot's bytes, masking only the score: include/flashinfer/attention/sparse_mla_sm120/common/kv_cache_io.cuh (io_bulk_gather_tile, io_gather_scales) and .../prefill_kernel.cuh (prefill_kv_entry_base, idx = (idx >= 0) ? idx : 0). With the score masked to -1e30 the probability is 0, but 0 * NaN = NaN in the value MMA, so whenever slot 0 of the cache the kernel is handed holds NaN-encoded fp8 / ue8m0 bytes (0x7F, 0xFF), every query row that carries a -1 in its indices comes out NaN. (The scalar RoPE-V path in xv_rope_mma already returned 0 for idx < 0 with a comment describing exactly this hazard; the bulk NOPE-V gather did not.)
2. sglang hands the kernel a scratch whose slot 0 is uninitialized. The SWA pool uses 256-token pages; the FlashInfer fast path wants 64-token pages, so _split_kv_pages_to_64 re-pages the pool into a persistent, grow-only scratch registered in get_resources().buffers. That buffer is allocated with torch.empty, and since the touched-page optimization only copies the source pages a step references, dst page 0 (slot 0) keeps whatever the caching allocator recycled into it unless source page 0 happens to be referenced. Every cold prompt has -1 padding in its top-k indices until the context exceeds the top-k width, so the NaN reaches every row of the prompt.

On upstream main the exposure is on the SWA split and depends on what the allocator recycled into the scratch and on whether the step references source page 0. A downstream variant of this module that also re-pages the ratio-1/2 extra cache through the same splitter hit it deterministically (sglang reserves full-pool page 0, so that scratch's page 0 was never written), which is how the 64/65 boundary was isolated; the fix is the same.

Kernel harness confirmation (direct _sparse_mla_sm120_paged_attention calls, single 64-page cache, H=16, indices with -1 padding):

kernel	T	slot 0 bytes	result	rows with a -1	NaN rows
FI prefill	65	0xFF	308,224 non-finite elements	43	43
FI prefill	65	random	4,128 non-finite	43	43
FI prefill	65	zeros	ok, max row rel err 0.034	43	0
FI decode	64	0xFF	non-finite	42	42
FI decode	64	zeros	ok, 0.056	42	0
Triton / torch ref	65	0xFF	ok, 0.007 / 0.001	43	0

The NaN rows are exactly the rows containing a -1. Every sglang stage of _flash_mla_sm120_prefill is byte-exact in isolation (256->64 split, index identity, gather + dequant through the 64-view, wrapper == direct call), and the kernels are numerically fine for every shape tested (T=8..1024 incl. 64/65/66, H=8..64, topk 128..2048, 1..6000 pages, causal / mixed lengths) once slot 0 is finite.

Realistic allocator recycling (a freed 0xFF-filled block is recycled by the caching allocator into the split scratch, then a production-shaped cold prompt runs):

T	path	torch.empty (current)	torch.zeros (this PR)
64	decode	ok 0.065	ok 0.065
65	prefill	65/65 rows NaN	ok 0.033, 0 NaN
128	prefill	127/128 rows NaN	ok 0.033, 0 NaN

The kernel-level reproduction script (poisons slot 0 of a small paged cache and calls the FlashInfer kernel directly) is available on request.

Modifications

One line in _split_kv_pages_to_64 (python/sglang/kernels/ops/attention/flash_mla_sm120.py): allocate the persistent split scratch with torch.zeros instead of torch.empty, so slot 0 is finite from the first use and stays so (it is only ever overwritten with real page-0 data). The buffer is grow-only, so the re-allocation path goes through the same line. The docstring and the decode-path comment that said untouched pages are "never read" are corrected: they are never addressed by a valid index, but slot 0 is gathered for every masked candidate by FlashInfer <= 0.6.18.

-            buf = torch.empty(
+            buf = torch.zeros(
                 num_dst_pages,
                 _BYTES_PER_DST_PAGE_PADDED,
                 dtype=torch.uint8,
                 device=dev,
             )

New test test/registered/unit/kernels/ops/attention/test_sm120_split_scratch.py (base-a-test-cpu, no GPU / torch / triton needed: the module is loaded by file path with a recording torch stub, the way unit/tools/ loads ci_register, and _split_kv_pages_to_64 is driven directly with its two Triton kernels replaced by launch recorders). It pins:

• the first allocation is torch.zeros((num_dst_pages, 37440), dtype=uint8, device) registered under flash_mla_sm120_split:<device>; the split kernel writes into a slice of that buffer with HAS_MASK=True; the returned view addresses it as 64-token pages;
• a same-size or smaller call reuses the buffer; a grow re-allocates with torch.zeros again; one scratch per device;
• the int8 page mask is zero_()ed on every call; a 64-token pool needs no scratch;
• AST: the only uint8 allocation in the module is that torch.zeros inside _split_kv_pages_to_64, no torch.empty uint8 buffer anywhere, and the allocation comment documents the slot-0 gather.

8 tests; 6 of them fail on the torch.empty variant (negative control run locally), all pass with the fix.

Relation to other changes

• feat(sm120): runtime KV row stride + canonical 528B GLM53_NOPE rows + NaN-safe masked gathers in sparse MLA flashinfer-ai/flashinfer#5075 (zero_row.cuh: masked candidates gather a shared zero row via mask_idx_past_len; merged 2026-09-11) fixes the kernel side. It is not in 0.6.18 (python/pyproject.toml pins flashinfer_python[cu13]==0.6.18) nor in any FlashInfer tag as of today (v0.6.18.post1 and the v0.7.0rc1 tag predate it). Either fix alone suffices; this PR makes sglang correct on the pinned FlashInfer and stays harmless once the pin moves past Support benchmarking with presets on multiple configuration combinations #5075 (the scratch can then be left zeroed or dropped with the splitter).
• fix(sm120): fall back to Triton per call when FlashInfer lacks the prefill shape #38969 ("fall back to Triton per call when FlashInfer lacks the prefill shape") is orthogonal: it adds _flashinfer_covers() envelope routing and leaves the torch.empty allocation in place, so covered shapes would still hand FlashInfer a poisoned slot 0.
• [SM120] DeepSeek-V4: DeepGEMM paged-MQA indexer +FP4 MoE+ page-split  #29927 (SM120 paged-MQA indexer + page split) introduced the scratch; [SM120] flash_mla: allocate the page-split buffer outside inference mode #35116 wrapped the same allocation in inference_mode(False) for CUDA-graph capture; [SM120] Use exact query-head widths for DeepSeek-V4 sparse MLA decode #36655 (exact query-head widths for decode) and [Bug] DeepSeek-V4 SM120 decode pads q to 64 heads for an SM90 constraint; removing the pad changes greedy output despite the kernel being bit-identical #39235 (SM120 decode q pad) touch the same dispatch and are unaffected. Fix int32 offset overflow in SM120 page-split kernel #34027 / test: pin the int32 page-offset overflow in _page_split_kernel #35833 (int32 offset overflow in _page_split_kernel) are about kernel indexing, not initialization.

Accuracy Tests

End to end (4x RTX PRO 6000, DeepSeek-V4.1-Flash, TP4, FlashInfer 0.6.18, FlashInfer >64-row path, no Triton fallback), same box patched vs unpatched:

	unpatched (torch.empty)	this PR (torch.zeros)
boot smoke: basic / structured / tools / vision	tools and vision FAIL 5x each	all pass
planted-code recall sweep, 23 cold prompts 36..32k tokens	5/23 (only prompts <= 64 tokens)	23/23, twice
262k-token cold prefill	wrong	correct
exact-fit oracle, 64 / 65 prompt tokens	4/4 correct / 0/4	correct / correct

With the fix the FlashInfer prefill path sits at ~3% max row relative error (BF16 P x V) against an fp32 reference, Triton at <1%; both far from garbage.

python test/registered/unit/kernels/ops/attention/test_sm120_split_scratch.py passes on CPU (8 tests).

Speed Tests and Profiling

No per-step cost: one memset per (re)allocation of the grow-only scratch (1.14 GiB at 2,097,152 SWA tokens in 256-token pages), i.e. once at the first split and once per grow. The per-step touched-page copy is unchanged.

Checklist

• Format your code according to the Format code with pre-commit. (ruff format, ruff check F401/F821/UP037, isort, codespell run locally on the two files)
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations. (none: no user-facing knob)
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

🤖 Generated with Claude Code

---

CI States

Latest PR Test (Base): ❌ Run #34747908972
Latest PR Test (Extra): ❌ Run #34747908830
Latest PR Test (AMD ROCm 10): ❌ Run #34747909096