[SM120] Add opt-in NVFP4 sparse-MLA KV cache format for DeepSeek-V4

[AliceChenyy]

Motivation

FlashInfer #4955 added native NVFP4 sparse-MLA kernels for SM120/SM121 next to the existing FP8 ones. This wires them into the DeepSeek-V4 attention path on SM120.

Opt-in through SGLANG_SM120_KV_CACHE_FORMAT=nvfp4. The default stays fp8 and no existing code path changes behavior.

Modifications

KV pool (deepseek_v4_memory_pool.py) — 384 B/token ABI (packed E2M1 nope + BF16 rope + E4M3 group scales) against FP8's 584 B. The NVFP4 ABI packs each page exactly, so the pool regroups the flat slot space into the page sizes the kernels accept (64 primary, 2 or 64 extra) with no per-page padding. Slot ids are unchanged by that regrouping, so writers and readers agree without remapping — and unlike the FP8 path, there is no 256→64 page-split copy on every attention call. Writes go through nvfp4_quantize_append_sparse_mla_cache.

Attention (flash_mla_sm120.py) — an NVFP4 branch that selects phase/CPB via plan_nvfp4_sparse_mla_sm120 (that planner skips calibration while a stream is capturing, so it is CUDA-graph safe) and memoizes the decision per shape.

Stores — three fused kernels only emit the FP8 ABI, so under NVFP4:

• the SWA store takes the existing bf16-intermediate path (SGLANG_DSV4_USE_BF16_KV_QUANT_SOURCE) and quantizes on store;
• the compressed c4/c128 store lands in a BF16 staging buffer through the existing bf16_store path and is quantized afterwards;
• the pre-quantized FP8 write path asserts rather than corrupting the cache silently.

Platform gate — is_nvfp4_kv_cache() raises at startup if the flag is set on anything other than SM120/SM121.

Reviewing deepseek_v4_backend.py: please use git diff -w. 13 of its 16 deleted lines are pure re-indentation from wrapping existing code in an else:; the remaining 3 are one statement black re-wrapped because the added indent pushes it past 88 columns. Semantic change there is +11/−3.

Performance

DeepSeek-V4-Flash, 4× RTX PRO 6000 Blackwell Server Edition (SM120, 96 GB), --moe-runner-backend marlin, --kv-cache-dtype fp8_e4m3. max_total_num_tokens pinned identically across each FP8/NVFP4 pair — NVFP4 stores 66% of FP8's bytes per token, so an unpinned NVFP4 arm gets a ~1.5× larger pool and its throughput is not comparable.

Gains depend on heads-per-rank, not on batch or sequence length

DSv4-Flash has 64 attention heads. Operator benchmark (FlashInfer's own, primary topk 128 = sliding_window, extra topk 512 = index_topk), NVFP4 vs FP8 prefill:

heads/rank	configuration	1024 tok	2048	4096	8192
16	attention TP4	−33.1%	−33.3%	−33.2%	−33.0%
32	attention TP2	−7.7%	−5.4%	−3.6%	−2.3%
64	DP-attention / PP / CP	+23.6%	+35.9%	+36.5%	+35.6%

Under attention TP the NVFP4 prefill kernel is slower than FP8. This is why the feature is opt-in and why the flag should only be enabled with attn_tp=1.

End-to-end, ISL 8192 (DP4: --tp 4 --enable-dp-attention --dp 4)

--disable-radix-cache, pinned to 3,990,784 tokens. Both arms measured against
the code in this branch, in the same container:

workload	FP8	NVFP4	delta
ISL8192 conc 8 × 64, throughput	14,811 tok/s	15,251 tok/s	+2.97%
↳ median / P99 TTFT	4383 / 5361 ms	4265 / 5209 ms	−2.7% / −2.8%
ISL8192 conc 32 × 256, throughput	14,811 tok/s	15,276 tok/s	+3.14%
↳ median / P99 TTFT	17,523 / 18,579 ms	17,037 / 17,919 ms	−2.8% / −3.6%

End-to-end, ISL 32768, three attn_tp=1 topologies

OSL=8, --num-prompts BS --max-concurrency BS --warmup-requests 0, --disable-radix-cache, --chunked-prefill-size 32768, every arm pinned to 3,864,064 tokens. Median of 5 samples per cell, measured against this branch:

topology	BS=1 TTFT	BS=1 throughput	BS=8 TTFT	BS=8 throughput
CP2×PP2	−5.5%	+5.7%	−5.9%	+5.9%
CP4 (--tp 4 --attn-cp-size 4 --enable-prefill-cp --cp-strategy interleave)	−3.6%	+3.9%	−3.7%	+4.2%
PP4 (--tp 1 --pp-size 4)	−7.1%	+7.6%	+0.8% ⚠️	+1.3% ⚠️

At BS=1 the gain tracks CP degree inversely, which is consistent with CP adding cross-rank communication that NVFP4 does not accelerate. That ordering reproduced across two independent passes.

⚠️ The PP4 BS=8 cell is bimodal under NVFP4 and should not be relied on. Its five samples split into two stable clusters roughly 9% apart, switching mid-run and not switching back:

PP4 BS=8 median TTFT (ms), 5 consecutive samples, one server boot
  FP8   : 10943  10899  10945  10921  10924     (spread 0.4%)
  NVFP4 : 10102  10134  11038  11014  11011     (two modes, 9% apart)

An earlier pass on the pre-fix build showed the same two modes in the opposite proportion (four fast, one slow), which is why a first 2-sample pair on this cell produced opposite-signed conclusions. FP8 is flat on the same shape, so this is specific to the NVFP4 path — plausibly the same tactic-selection issue as the 64-query-token operator regression noted under Limitations. Every other cell reproduced within 0.6 pp across the two passes.

For reference, FlashInfer #4955 reports +9.02% token throughput for a vLLM PP=4 / ISL8192 / OSL1 run.

Accuracy

GSM8K, 5-shot, DP4:

n	FP8	NVFP4
1000 (stderr ≈ 0.54 pp)	0.967	0.965

No measurable loss (0.3σ). An earlier n=200 pass showed a 0.5–1.0 pp dip; at n=1000 that does not hold up, so it was noise.

Limitations

• No fused NVFP4 store kernel. The NVFP4 path gives up the fused norm+rope+quantize+store kernel and pays a BF16 intermediate plus a separate quantize kernel per layer, plus a staging-buffer clear for the compressed cache. The numbers above are therefore a floor; a fused NVFP4 store should recover part of the gap to [Fix] error when processing batch prompts for VLM #4955's reported figure.
• Requires FlashInfer main (≥ [Fix] error when processing batch prompts for VLM #4955). If flashinfer-jit-cache is installed from a wheel predating that PR, its prebuilt FP8 module has the old 12-argument signature and the FP8 path breaks too; uninstall it so the kernels JIT from source.
• The operator benchmark shows a reproducible regression at exactly 64 query tokens across all head counts, which looks like a planner tactic choice. Not investigated; reported here for visibility.

Checklist

• black, isort clean; ruff introduces no new findings
• Format defaults to fp8; FP8 behavior unchanged
• ISL 8192 numbers and GSM8K re-measured against this branch, both arms in the same container
• ISL 32768 topology table re-run against this branch, 5 samples per cell, all six arms on one KV budget
• Regression tests added in test/srt/mem_cache/test_dsv4_nvfp4_cache.py (14 cases). Each one pins a defect that reached a running server during development: staging-buffer pointer stability, bucketing, per-call zeroing, int64 out_loc, uint8 coercion in nvfp4_cache_view, and the platform gate. All of that runs on CPU; only the slot-addressing round-trip needs the kernels and is skipped off SM120 (13 passed, 1 skipped with no GPU visible, 14 passed on an RTX PRO 6000). Verified by mutation: restoring the grow-in-place staging buffer turns two of them red with staging storage moved; a replayed CUDA graph would write through a freed pointer. Not registered in run_suite.py — sglang CI has B200 runners but no SM120, so there is nowhere for the GPU case to run; happy to wire it in if a suitable suite exists.

Two bugs found in self-review and fixed in this branch, noted for reviewers since they are subtle. They cost roughly 0.9 pp of throughput (an earlier build measured +3.9/+4.0% where this one measures +3.0/+3.1%); the FP8 baseline reproduced to within 0.3% across the container rebuild that separates those runs, so the difference is more likely real than environmental, though each NVFP4 configuration is a single sample.

1. The compressed-cache staging buffer was reallocated when a prefill needed more rows than the decode CUDA graph was captured with, leaving the replayed graph writing through a freed pointer. It happened to produce correct results because the captured store and the captured quantize shared the same stale pointer.
2. Keying that buffer on the exact row count made the cache unbounded, since prefill row counts follow each chunk's token count. It is now bucketed to a power of two and sliced, which keeps the backing pointer stable and the buffer count bounded.

🤖 Generated with Claude Code