[Perf][DSA] Grouped sparse prefill with the FlashMLA head-group-mask kernel - #6
JaredforReal wants to merge 1 commit into
Conversation
…kernel The per-token sparse MLA prefill kernel re-gathers every selected KV row for every query token (2048 x 1152 B per token); with TP4 all four ranks gather the same rows, so prefill attention is HBM-gather bound (~4.8 ms/layer for a 16k chunk, ~45% of the prefill step on GLM-5.3-NVFP4). Consecutive queries select heavily overlapping top-k sets: on real prompts the union of 8 adjacent tokens' top-2048 is ~2.4k keys. Present 128 // num_heads consecutive tokens of a request as 128 pseudo-heads that attend over the union of their top-k rows with FlashMLA's SM100 sparse prefill kernel, using its new `head_group_mask` argument to restrict each token's heads to its own top-k, which reproduces per-token sparse attention exactly (kernel test vs. an exact reference: rel err 2e-3). Three Triton kernels build the sorted union list and the membership bitmask from the indexer top-k (0.56 ms per 16k tokens). Opt-in via VLLM_DSA_GROUPED_SPARSE_PREFILL=1 (bf16 KV, no DCP); VLLM_DSA_GROUPED_SKIP_MASKED_MHA=1 additionally routes every prefill longer than the top-k budget through it instead of the masked-MHA path. Requires a FlashMLA build with `head_group_mask` support. Attention kernel time on a real 32k prompt (layer 10 top-k, 21k queries): 14.25 ms -> 2.86 ms. End to end on GLM-5.3-NVFP4 TP4 / 4x GB300, on top of the fused_q change: P16k TTFT -10%, P32k -21..26%, P64k -23%, 32k-context c=16 +27% tok/s (-36% / -32% / +42% vs. main); prompt-logprob differences vs. the per-token path are within the baseline's run-to-run noise. Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Jared Wen <jaredwen@inferact.ai>
There was a problem hiding this comment.
🟡 Changes recommended
There is a confirmed crash path when the env var is enabled without FlashMLA being available, plus test/perf issues that should be addressed before landing.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
Adds an experimental “grouped sparse prefill” path for DSA models that batches consecutive prefill tokens to share a union of top‑k rows and uses FlashMLA’s head-group-mask support to preserve per-token sparsity semantics. This is gated by env vars and integrates into the FlashInfer sparse MLA backend selection logic.
Changes:
- Add a grouped-prefill execution path in
FlashInferMLASparseImpl.forward_mqathat routes prefill tokens through FlashMLA with a head-group mask. - Introduce Triton kernels + reference implementation to build per-group union index lists and per-token membership masks.
- Add a CUDA kernel test validating Triton metadata builder vs the reference, plus env-var behavior to skip masked-MHA when enabled.
File summaries
| File | Description |
|---|---|
| vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py | Adds env-var gated grouped-prefill path that calls FlashMLA sparse kernel with head-group masks. |
| vllm/models/deepseek_v32/nvidia/ops/grouped_sparse_prefill.py | New Triton implementation for grouped union + mask metadata generation. |
| vllm/model_executor/layers/attention/mla_attention.py | Adds env-var switch to disable masked-MHA when grouped sparse prefill is enabled. |
| tests/kernels/test_grouped_sparse_prefill.py | New CUDA test comparing Triton metadata builder to a torch reference. |
Review details
- Files reviewed: 4/4 changed files
- Comments generated: 3
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| union_pad = triton.cdiv(group * num_topk, UNION_BLOCK) * UNION_BLOCK | ||
| words = triton.cdiv(max_seq_len, 32) | ||
| nblk = union_pad // UNION_BLOCK | ||
|
|
| def _use_grouped_prefill(self, attn_metadata: FlashInferMLASparseMetadata) -> bool: | ||
| return ( | ||
| _GROUPED_SPARSE_PREFILL | ||
| and attn_metadata.num_prefills > 0 | ||
| and attn_metadata.prefill is not None | ||
| and self.dcp_world_size <= 1 | ||
| and not is_quantized_kv_cache(self.kv_cache_dtype) | ||
| and not self.need_to_return_lse_for_decode | ||
| and _GROUPED_PSEUDO_HEADS % self.num_heads == 0 | ||
| and _GROUPED_PSEUDO_HEADS // self.num_heads <= 8 | ||
| ) |
| scores = torch.rand(num_tokens, max_seq_len, device="cuda") | ||
| idx = scores.argsort(dim=1)[:, :topk].int().contiguous() |
|
GLM-5.3-Flash (zai-org, FP8, TP4 on 4x GB300) results with this branch cherry-picked onto the
No code change was needed for GLM: the FlashMLA sm100 sparse prefill kernel and the head-group-mask patch are templated on D_QK (512 here), and the union builder already handles the kpool -1 padding and the wider (2176-column) top-k buffer. |
Summary
Grouped sparse prefill for DSA models (GLM-5.3 / DeepSeek V3.2) on SM100, stacked on #5. Off by default (
VLLM_DSA_GROUPED_SPARSE_PREFILL=1), requires a FlashMLA build with thehead_group_maskargument (JaredforReal/FlashMLA#1).Why
Profiling GLM-5.3-NVFP4 prefill (16k-token chunks) showed the FlashInfer per-token sparse MLA kernel at ~4.8 ms/layer = 45% of the step: every query token gathers its 2048 × 1152 B of KV, and with TP4 all four ranks gather the same rows — it is HBM-gather bound, not compute bound. Measured on real prompts, 8 adjacent tokens' top-2048 sets have a union of only ~2.4k keys (
tools/topk_locality.py).How
nvidia/ops/grouped_sparse_prefill.py: three Triton kernels build, per group of128 // num_headsconsecutive tokens of a request, the sorted union of their top-k (bitmap → prefix-sum compaction) and a per-token membership bitmask in FlashMLA'shead_group_masklayout. 0.56 ms per 16k tokens incl. Python. Torch reference +tests/kernels/test_grouped_sparse_prefill.py.flashinfer_mla_sparse.py:forward_mqasplits decode tokens (unchanged FlashInfer path) from prefill tokens, which go through_forward_grouped_prefill: union → physical rows viatriton_convert_req_index_to_global_index,qgathered into[groups, 128, 576],flash_mla_sparse_fwd(..., topk_length, head_group_mask), padding rows dropped.mla_attention.py:VLLM_DSA_GROUPED_SKIP_MASKED_MHA=1keeps the dense-MHA shortcut for prompts within the top-k budget but routes everything longer through the grouped kernel instead of masked MHA (masked MHA loses to it at every length).Results (GLM-5.3-NVFP4, TP4, 4× GB300, bf16 KV, prefix caching off; same script, same window)
Kernel only (real layer-10 top-k, 21k queries, 32k KV): FlashInfer per-token 14.25 ms → FlashMLA grouped 2.86 ms (5.0×).
main: 32k TTFT −36%, 64k −32%, 32k-ctx c=16 throughput +42%.Correctness
qcarries the prefill tokens).Tests
pytest tests/kernels/test_grouped_sparse_prefill.py(6 passed),tests/kernels/test_fused_deepseek_v32_norm_rope.py -k fused_q(30 passed). E2E scripts and logs under~/notes/glm53-nvfp4-perf/(tools/e2e_grouped.sh,tools/logprob_ab.sh).Open points for review
vllm/third_party/flashmla/flash_mla_interface.pyis regenerated from it at build time.