[Kernel] MXFP4 indexer cache for GLM-5.2 / DSA (glm_moe_dsa)

[nsd9696]

[Kernel] MXFP4 indexer cache for GLM-5.2 / DSA (glm_moe_dsa)

What

Adds MXFP4 (4-bit) support for the DSA lightning-indexer K-cache and wires it into
GLM-5.2 (GlmMoeDsaForCausalLM, deepseek_v2 Indexer). Gated by the existing
--attention-config '{"use_fp4_indexer_cache": true}' (default off, behavior unchanged).

Upstream shipped the FP4 indexer read path (mqa_logits) for DeepSeek-V4 in #40860, but
the insert path is still blocked — assert not use_fp4_cache, "Unfused FP4 Insert is not supported yet" — and indexer_k_quant_and_cache is FP8-only. This PR fills that gap.

Changes (3 files, +122 / −3)

• csrc/.../cache_kernels.cu (new kernel): indexer_k_quant_and_cache_mxfp4_kernel
  writes the indexer K-cache as MXFP4 (E2M1 values + UE8M0 block-32 scales, software E2M1
  encode for sm_100a). Dispatched from indexer_k_quant_and_cache when scale_fmt=="mxfp4".
• sparse_attn_indexer.py: drop the "not supported yet" assert; pass "mxfp4" scale_fmt.
• deepseek_v2.py (Indexer): allocate the fp4-shaped K-cache, MXFP4-quantize the
  indexer query, disable the fused FP8 indexer-q path.

Benefit

Indexer K-cache: 132 → 68 B/token/head (~1.94×). Measured on 4×B200 (GLM-5.2-NVFP4, TP4):

metric	FP8 indexer	MXFP4 indexer	Δ
GPU KV cache size	1,143,808 tok	1,184,192 tok	+3.5%
max concurrency @131k	8.73×	9.03×	+3.4%
TTFT / TPOT (conc1 prefill 8K–122K; conc16 grid; conc48–64 long-ctx)	—	—	neutral

It's a memory-efficiency feature (more KV headroom at long context), not a latency win —
the indexer is a small fraction of prefill vs MLA attention.

Accuracy

passkey 8K / 32K / 64K = 3/3 · 3/3 · 3/3, identical to FP8. No regression.

Testing

• Kernel + wiring validated end-to-end on 4×B200 (SM100) serving GLM-5.2-NVFP4: fp4 indexer
  cache active (KV size 1.18M tok), coherent generation, passkey 3/3.
• The use_fp4_indexer_cache=true flag was exercised on current-main Python (fp4 K-cache
  allocated, 9.09× concurrency). The modified C++ TU (cache_kernels.cu = main + this diff)
  compiles clean against main's headers with CUDA 13.0 / sm_100 (exit 0, no warnings;
  see the compile-verification comment below). Full multi-arch build + link across all TUs is
  still what upstream CI does — a ready label would run it.
• Requires SM100 datacenter Blackwell (B200/GB200) — same constraint as the existing
  use_fp4_indexer_cache backend assert.

Notes

• MXFP4 q-quant is done in PyTorch (indexer q is tiny vs attention); a fused q-rope-mxfp4
  kernel could follow.
• GLM's index_n_heads=32 drives the paged MQA-logits reader at N=16; the vendored
  deep_gemm needs N∈{8,16} TMEM loaders (a 1-line assert/loader extension, tracked
  separately from this vLLM change).

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[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.

[nsd9696]

Live validation on 4×B200 (SM100), GLM-5.2-NVFP4, TP4 + EP

Posting the end-to-end evidence behind the "passkey 3/3 / latency-neutral" claims, in case it helps a maintainer decide on a ready label.

FP4 indexer cache is active (server boot log):

GPU KV cache size: 1,184,192 tokens        # fp4 indexer cache (fp8 baseline: 1,143,808)
Maximum concurrency for 131,072 tokens per request: 9.03x   # fp8 baseline: 8.73x

The indexer K-cache byte width drops from 132 → 68 B/token/head, which is the +3.5% KV headroom.

Passkey retrieval — 9/9 HIT (temperature 0, single needle, 3 trials per length):

context (~prompt tok)	trials	result
8K (≈11.5K ptok)	3	HIT / HIT / HIT
32K (≈47.5K ptok)	3	HIT / HIT / HIT
64K (≈97K ptok)	3	HIT / HIT / HIT

Identical to the FP8 indexer path — no accuracy regression.

Latency: neutral vs FP8 across conc1 prefill (8K–122K), a conc16 grid, and conc48–64 long-context. The indexer is a small fraction of prefill vs MLA attention, so this is a memory-efficiency feature (more KV headroom at long context), not a latency win.

Provenance of the kernel under test. The indexer_k_quant_and_cache_mxfp4_kernel in this PR is byte-identical to the one compiled into the binary that produced the numbers above:

served binary  sha256 = 2b27f50…
built .so      sha256 = 2b27f50…   (same file; source cache_kernels.cu → compiled 1 min later)
PR kernel      == that source's kernel   (diff 0)

So the C++ in this PR is exactly what was validated end-to-end; the block-split value/scale layout matches the existing FP8 insert path + the shared cp_gather / paged MQA-logits reader.

[nsd9696]

Update: the C++ change now compiles clean against main

The earlier caveat ("not compile-verified against main") is resolved for the kernel change. I compiled the modified translation unit — csrc/libtorch_stable/cache_kernels.cu at this PR's HEAD (i.e. main + this diff) — against main's own csrc headers with the target toolchain:

• Toolchain: CUDA 13.0 (nvcc V13.0.88), C++20, project RelWithDebInfo flags (--expt-relaxed-constexpr --expt-extended-lambda -O2 -DENABLE_FP8 -DENABLE_NVFP4_SM100=1 …), headers from main's csrc + Torch stable ABI + CUTLASS.
• Target: -gencode arch=compute_100,code=sm_100 (B200/SM100; the kernel itself is arch-generic).
• Result: exit 0, zero warnings, zero errors, cache_kernels.cu.o produced.

Symbols confirmed in the object (so the new kernel is emitted and actually referenced by the dispatch, not dead-stripped):

T indexer_k_quant_and_cache_mxfp4_kernel(__nv_bfloat16 const*, unsigned char*, long const*, int, int, int)
T __device_stub__..._mxfp4_kernel...          # dispatch call site present
T indexer_k_quant_and_cache(torch::stable::Tensor&, ...)
$ cuobjdump cache_kernels.cu.o | grep arch
arch = sm_100                                   # SASS generated for target

This only covers the single file this PR touches (the other two changes are pure Python). The full multi-arch build + link across all TUs is still what upstream CI does — which is why I'd still appreciate a ready label to run it — but the specific "does the new C++ compile against current main" question is now answered: yes, clean.