[None][feat] Enable DeepSeek-V4 and DSA (DeepSeek-V3.2/GLM) serving on SM120 via FlashInfer sparse-MLA

[lucifer1004]

SM120 compiles out the trtllm-gen sparse-MLA path in the C++ attention op for both phases (attentionOp.h excludes SM120 from mUseTllmGen; generation asserts at attentionOp.cpp:3156), so neither DeepSeek-V4 nor DSA models (DeepSeek-V3.2, GLM glm_moe_dsa) can serve on SM120/SM121. This PR routes both families through flashinfer 0.6.14's SM120 sparse-MLA kernels (footer-scale cache for DSv4, inline-scale for DSA). All-Python change; no C++ rebuild required. Four commits: DSv4 backend, DSv4 reuse policy, DSA RoPE fix, DSA backend — each separable.

DeepSeek-V4 (commits 1–2)

• DeepseekV4FlashInferAttention reuses the TRTLLM backend's metadata, indexer, compressor, and local-to-global index conversion; one flashinfer call serves context and generation (decode fast path ≤64 query tokens, prefill orchestrator above).
• SWA and COMPRESS pools switch to the kernels' footer-scale page layout (per token: 448 B fp8-e4m3 nope + 128 B bf16 rope + a UE8M0 tile-scale footer row); requires kv_cache_config.tokens_per_block=256 so compressed block sizes (64, 2) land on the instantiated extra page sizes.
• The layout switch derives from model_config.attn_backend on both the cache-manager and backend sides; the fused C++ compressor scatter is replaced by an equivalent Python postprocess + Triton footer-scale scatter; the SWA append moves into the backend; RoPE runs in mla.py (support_fused_rope False), with DSv4's fused-rope and fused-FP8-Q paths gated on rope_fusion.
• The local-to-global index kernel gains an optional split-segment output (combined table unchanged for the C++ consumer); decode split-K scratch is explicitly initialized; all pool byte offsets are 64-bit (int32 slot-id byte products overflow past a 2 GiB pool, reached at ~1M-token contexts).
• The reuse-policy commit keeps KV block reuse enabled for DeepSeek-V4 under this backend (the FLASHINFER reuse workaround targets the dense wrapper); it depends on [None][fix] Derive DSA indexer K-cache block stride at use time #16223 and is separable if reviewers prefer to take it independently.

DSA family — DeepSeek-V3.2 / GLM (commits 3–4)

• DSAFlashInferAttention reuses the TRTLLM DSA backend's metadata, indexer (including GLM-5.2's cross-layer top-k sharing), and local-to-global index conversion; one flashinfer call serves context and generation, with decode split-K scratch explicitly initialized. Backend selection rejects non-SM120/121 with a clear error.
• The main latent pool switches to the kernels' inline-scale page layout (FlashMLA ABI: 512 B fp8-e4m3 nope in 4×128 tiles + 4 inline fp32 scales + 128 B bf16 rope per token; GLM-style arbitrary-fp32 scale mode), 64-token pages. The Triton writer is byte-identical to SGLang's quantize_k_cache. BF16- and FP8-declared KV checkpoints land in the same physical layout.
• The append moves into the backend: write slots come from the same convert_req_index_to_global op the top-k conversion uses, fed each token's own absolute position (refreshed per step alongside the indexer slot mappings, overlap-scheduler-safe).
• The softmax scale uses the per-head qk width from mla_params rather than the absorbed latent width the backend head_dim carries — they coincide for DSv4 (512) but differ for DSA geometries (GLM-5.2: 256 vs 576).
• RoPE runs in Python here too, which surfaced two latent defects in the never-before-exercised DSA + Python-RoPE combination, fixed in commit 3 entirely within forward_dsa_attn's phase split: the generation phase pre-rotated q_pe/k_pe that forward_absorption_generation rotates again (the pattern is only correct for DeepSeek-V4, which skips _mla_gen_rope), and the context phase never propagated the rotated k_pe into the latent cache. Shared absorption functions and fused-RoPE backends are untouched.
• The FLASHINFER reuse workaround exemption extends to the DSA sparse backend on the same grounds as DSv4.

Validation (RTX PRO 6000 Blackwell, SM120; closed-loop serve+eval+bench with verified cleanup)

DeepSeek-V4 (4 GPUs):

• gsm8k exact_match 0.945–0.957 at tp2/ep2 and tp4/ep4 (threshold 0.8), serving bench green across concurrency ramps (reuse off and on).
• With reuse enabled: the previously-crashing 8k-prompt concurrency ramp is green through c=32; prefix-resume semantics verified across a mid-block partial-reuse boundary.
• Long-context ladder on the shipped tree (tp4, chunked prefill): 128k/512k/1M all green — three planted-fact retrieval correct at every rung, cold prefill 6.8–8.1k tok/s, prefix-reuse repeat up to 52.7× TTFT gain at 1M.
• The packed-page writer is byte-identical to SGLang's reference packer; kernel numerics validated against SGLang's torch reference (min cosine ≥ 0.9994).

DSA family, GLM-5.2-NVFP4 and GLM-5.1-NVFP4 (8 GPUs, tp8/ep8):

• gsm8k exact_match 0.9568 (threshold 0.8).
• Serving bench (8k-in/1k-out): 289 output tok/s peak at concurrency 16, 21.7 ms single-user ITL (first bring-up; chunked prefill off).
• 92-request tool-use campaign (2k–10.8k-token prompts, glm47 tool parser, greedy, block reuse and partial reuse on, including truncation sweeps with >96% partial prefix reuse): all clean.
• Inline-scale writer byte-identical to SGLang's packer; kernel numerics vs a torch reference over the dequantized cache min cosine ≥ 0.998 across decode/prefill dispatch (the arbitrary-fp32 mode quantizes softmax weights to fp8 base+residual by design).

Depends at runtime on #16217 (DeepGEMM SM120 indexer kernels) and flashinfer ≥ 0.6.14; the reuse commit depends on #16223. Builds cleanly with #16188/#16215/#16216 for externally managed build environments.

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features

  • Added FlashInfer-based sparse attention support for DeepSeek-V4.
  • Introduced an optimized footer-scale KV-cache layout with quantization and dequantization support.
  • Added configurable split handling for compressed attention indices.
  • Improved KV-cache sizing and buffer management for supported attention backends.

• Bug Fixes

  • Improved rotary embedding and latent-cache handling across attention execution paths.
  • Preserved compatible block-reuse behavior for DeepSeek-V4 FlashInfer execution.

[coderabbitai]

Caution

Review failed

An error occurred during the review process. Please try again later.

📝 Walkthrough

Walkthrough

DeepSeek-V4 adds a FlashInfer sparse MLA backend using footer-scale KV pages, Python compression scatter, split sparse indices, updated RoPE handling, backend-specific cache sizing, and executor routing.

Changes

DeepSeek-V4 FlashInfer integration

Layer / File(s)	Summary
Footer-scale layout and cache accounting tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.py, tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py	Cache sizing and token-byte calculations now support native and footer-scale layouts, with footer-scale buffers represented as UINT8 containers.
Footer-scale quantization and compression tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/footer_scale_kv.py, tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/compressor.py	Adds footer-scale quantization, dequantization, and Python postprocess/scatter support for compressed KV rows.
Split sparse-index output tensorrt_llm/_torch/attention_backend/sparse/kernel.py, tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/deepseek_v4.py	Sparse index conversion optionally returns sliding-window and compressed index regions separately while preserving the combined default behavior.
FlashInfer sparse attention path tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/flashinfer.py, tensorrt_llm/_torch/attention_backend/sparse/utils.py	Introduces the DeepSeek-V4 FlashInfer attention implementation, footer-scale pool handling, KV appends, and dispatch registration.
RoPE and executor routing tensorrt_llm/_torch/modules/mla.py, tensorrt_llm/_torch/pyexecutor/_util.py, tensorrt_llm/_torch/pyexecutor/py_executor_creator.py	Updates external RoPE cache writes and defers FlashInfer block-reuse decisions until sparse-attention routing is known.

Estimated code review effort: 5 (Critical) | ~120 minutes

Suggested reviewers: brb-nv, yuxianq, yihwang-nv, pengbowang-nv, dongxuy04, JunyiXu-nv, juney-nvidia, mzweilz

Sequence Diagram(s)

sequenceDiagram
  participant Executor
  participant DeepseekV4FlashInferAttention
  participant CacheManager
  participant Compressor
  participant FlashInferMLA
  Executor->>CacheManager: create footer-scale cache layout
  Executor->>DeepseekV4FlashInferAttention: route DeepSeek-V4 attention
  DeepseekV4FlashInferAttention->>Compressor: postprocess and scatter compressed KV
  DeepseekV4FlashInferAttention->>CacheManager: append SWA KV rows
  DeepseekV4FlashInferAttention->>FlashInferMLA: run paged sparse MLA attention

Loading

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[lucifer1004]

Scope update: pushed two additional commits extending this enablement to the DSA family (DeepSeek-V3.2 / GLM glm_moe_dsa) — [fix] for two latent defects in the DSA + external-RoPE combination (double rotation in the generation phase split, missing latent-cache rope sync in context; unreachable before a non-fused DSA backend existed), and [feat] for DSAFlashInferAttention on the inline-scale (FlashMLA-ABI) SM120 kernels. Title and description updated accordingly; all four commits are separable. Validated with GLM-5.2-NVFP4 and GLM-5.1-NVFP4 at tp8/ep8: gsm8k 0.9568, serving bench and a 92-request tool-use campaign with block reuse on, all clean — details in the description.