[Perf][DSA] Token-blocked fused_q kernel and packed bf16 MQA query

[JaredforReal]

Summary

Two related changes on the DSA (DeepSeek V3.2 / GLM-5.x) attention prefix path in vllm/models/deepseek_v32/:

1. _fused_q_kernel rewrite (common/kernels.py): grid (cdiv(tokens, 16), tasks) — one program handles 16 tokens × one head of one task (MQA q_pe RoPE, indexer RoPE+UE8M0 quant, or the ql_nope pack). The old grid (tokens, 3, 32) launched one single-warp program per element group and, on the bf16-query path, two thirds of the programs returned immediately. 16k tokens on GB300: 1582 µs → 152 µs.
2. Packed bf16 MQA query (attention.py, fused_q(..., q_out=)): the absorbed q_nope @ W_UK_T bmm writes directly into a [tokens, heads, kv_lora_rank + rope] buffer (torch.bmm(out=strided view), verified bit-equal and not slower) and fused_q RoPEs q_pe into the tail, so forward_mqa gets one contiguous tensor and the per-layer torch.cat in the sparse backends (518–645 µs per layer at 16k tokens, ~5 µs per decode layer) disappears. The (ql_nope, q_pe) tuple path is no longer used by this model; backends still accept it.

Results (GLM-5.3-NVFP4, TP4, 4× GB300, FLASHINFER_MLA_SPARSE, bf16 KV, prefix caching off, same script / same window vs main)

case	main	this PR
1k/512 c=1 median TPOT	7.83 ms	7.69 ms
1k/512 c=64 TPOT / tok/s	22.01 ms / 2657	21.27 ms / 2764
1k/512 c=256 TPOT / tok/s	45.57 ms / 5031	42.40 ms / 5654
32k prefill c=8 mean TTFT	12043 ms	10850 ms (−9.9%)
64k prefill c=2 mean TTFT	6286 ms	5552 ms (−11.7%)
32k ctx / 256 out c=16 tok/s / TPOT	145.5 / 82.1 ms	163.2 / 74.7 ms

Tests

• pytest tests/kernels/test_fused_deepseek_v32_norm_rope.py -k fused_q: 30 passed (adds test_fused_q_bf16_query_packed, incl. the aliasing case where ql_nope already lives in the buffer).
• Microbenchmark: notes/glm53-nvfp4-perf/bench/fused_q_bench.py.
• Output sanity on a 30k/60k real document (prompt logprobs) unchanged within run-to-run noise.

Notes

Not a duplicate of any open upstream PR (gh pr list --search "fused_q", "concat_mla_q" — none touch this). AI assistance (Claude Code) was used; the submitter reviewed every line. Model quality eval on gsm8k/aime is still to be run before an upstream submission.

[Copilot]

🟡 Changes recommended

The new packed-tensor mqa_q path is incompatible with sparse MLA backends that require a split (q_nope, q_rope) tuple (e.g. FLASH_ATTN_MLA_SPARSE, FLASHINFER_MLA_SPARSE_SM90), causing runtime NotImplementedError.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

This PR optimizes the DeepSeek V3.2 / GLM-5.x sparse-MLA prefix path by (1) restructuring the Triton fused_q kernel launch into token-blocked, per-head “task” programs, and (2) enabling a packed bf16 MQA query buffer to eliminate per-layer concatenations on the bf16-query path.

Changes:

• Rewrite _fused_q_kernel to process BLOCK_T tokens per program and unify MQA RoPE, indexer RoPE+quant, and optional NoPE packing under a single (token_block, task) grid.
• Extend fused_q(..., q_out=...) to write RoPE’d q_pe into the tail of a provided bf16 packed query buffer and optionally copy/pack ql_nope into the front.
• Update DeepSeek v3.2 attention to use the packed bf16 query path, and add tests covering q_out packing (including front-aliasing).

File summaries

File	Description
vllm/models/deepseek_v32/common/kernels.py	Reworks Triton fused-q kernel scheduling and adds optional packed bf16 query output via q_out.
vllm/models/deepseek_v32/attention.py	Uses a single packed bf16 query buffer to avoid torch.cat on bf16-query sparse decode.
tests/kernels/test_fused_deepseek_v32_norm_rope.py	Adds coverage for packed bf16 q_out behavior, including aliasing of the front slice.

Review details

• Files reviewed: 3/3 changed files
• Comments generated: 1
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.