[Perf][GLM-5.3-Flash] Dense/masked-MHA sparse prefill for the NoPE (256, 0, 256) layout + skip the NoPE K concat - #55738
Merged
ZJY0516 merged 10 commits intoSep 14, 2026
Conversation
…tention prefill backend GLM-5.3-Flash (qk_nope 256, qk_rope 0, v 256) was not in FlashAttnPrefillBackend.supports_mla_dimensions, so sparse MLA fell back to "No MLA prefill backend supports this model" and every prefill token went through the per-token top-k MQA kernel. The kernels for qk_head_dim 256 / v_head_dim 256 are the ones already used for the (192, 64, 256) layout. Same 7 lines as in the SM90 wiring PR (vllm-project#55385); kept as a separate commit so it can be dropped once that lands. Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Jared Wen <jaredwen@inferact.ai>
…4, 512, 256, 0, 256) layout wire_sm90_kernels registers the (256, 0, 256) MLA dimensions with the FlashAttention prefill backend, which gives GLM-5.3-Flash the dense-MHA prefill path for sequences up to index_topk (2048). Sequences above that still went through the per-token top-k MQA kernel on SM100, which is KV-gather bound (5.9 ms/layer per 16k-token chunk, ~4.8 ms of it pure HBM traffic), because the masked-MHA allow-list only knew the (128, 512, 128, 64, 128) and (64, 512, 192, 64, 256) layouts. The masked kernel path is identical for qk_head_dim 256 / v 256, so add the GLM layout; the existing FLASHINFER_MLA_SPARSE TP4/TP8 thresholds (masked MHA up to 36k / 64k tokens) apply. Prefill TTFT on 4x GB300 (TP4, together with the FlashKDA and decode cleanup PRs): 8x2048 462 -> 363 ms, 2x8192 382 -> 317 ms, 32k/256 c16 TTFT 3876 -> 3425 ms. First use of a new sequence-length bucket JIT compiles the FA4 mask_mod variant (~7 s); the FA4 MLA prefill warmup does not cover mask_mod yet. Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Jared Wen <jaredwen@inferact.ai>
With qk_rope_head_dim == 0 the concatenated K is just k_nope; return it instead of allocating and copying a same-sized tensor (134 MB per layer for a 16k-token chunk on GLM-5.3-Flash). Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Jared Wen <jaredwen@inferact.ai>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
Included review availability: Your plan provides up to 10 included reviews per hour; 6 remain after this review. WalkthroughMLA execution now handles zero-width RoPE tensors without concatenation. Masked MHA and FlashAttention prefill support the GLM-5.x dimensions ChangesGLM-5.x NoPE MLA support
Estimated code review effort: 2 (Simple) | ~10 minutes Merge Risk: ⚪ Minimal · up to This change enables optimized FlashAttention prefill support for GLM-5.x NoPE MLA dimensions. No concrete merge-blocking risk remains in the available evidence. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
JaredforReal
marked this pull request as ready for review
September 7, 2026 13:15
JaredforReal
requested review from
LucasWilkinson,
MatthewBonanni and
pavanimajety
as code owners
September 7, 2026 13:15
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
ZJY0516
approved these changes
Sep 7, 2026
Member
|
/ci run |
|
✅ @JaredforReal, CI is now available for this PR.
|
|
✅ Triggered Buildkite CI #87552 for commit |
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
The new hard-coded support/allow-list entries and NoPE concat fast-path should be accompanied by focused regression tests to prevent future backend-selection and correctness regressions.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Pull request overview
This PR improves GLM-5.3-Flash prefill performance by routing the NoPE MLA layout (qk_nope=256, qk_rope=0, v=256) onto faster dense/masked-MHA paths when applicable, and by avoiding an unnecessary key concatenation copy for NoPE models.
Changes:
- Register the
(256, 0, 256)MLA dimensions as supported by the FlashAttention prefill backend. - Extend the masked-MHA allow-list to recognize GLM-5.3-Flash’s NoPE layout on SM100-family GPUs.
- Skip allocating/copying concatenated K when the RoPE component is empty (
pe_dim == 0), returningk_nopedirectly.
File summaries
| File | Description |
|---|---|
vllm/v1/attention/backends/mla/prefill/flash_attn.py |
Adds (256, 0, 256) to FlashAttention MLA dimension support so NoPE layouts can use the dense prefill path. |
vllm/model_executor/layers/attention/sparse_mla_attention.py |
Adds GLM-5.3-Flash NoPE dims to the masked-MHA allow-list so mid/long prefills can use FA4 masked MHA instead of per-token sparse MQA. |
vllm/model_executor/layers/attention/mla_attention.py |
Avoids K concat allocation/copy for NoPE by returning k_nope when k_pe is empty. |
Review details
- Files reviewed: 3/3 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.
Comment on lines
+2792
to
+2795
| if k_pe.shape[-1] == 0: | ||
| # NoPE MLA: nothing to append, so no copy either. | ||
| return k_nope | ||
|
|
Comment on lines
98
to
103
| if model_dims not in ( | ||
| (128, 512, 128, 64, 128), | ||
| (64, 512, 192, 64, 256), | ||
| # GLM-5.3-Flash: NoPE, qk_head_dim 256 == the (192, 64, 256) kernel. | ||
| (64, 512, 256, 0, 256), | ||
| ): |
Comment on lines
+330
to
+336
| # GLM-5.x NoPE layout: qk_head_dim 256 + 0 and v_head_dim 256 run | ||
| # the same kernels as the (192, 64, 256) DeepSeek-V3.2 layout. | ||
| MLADimensions( | ||
| qk_nope_head_dim=256, | ||
| qk_rope_head_dim=0, | ||
| v_head_dim=256, | ||
| ), |
Signed-off-by: Jared Wen <w13431838023@gmail.com>
Member
|
@JaredforReal please add some unit tests, the coplilot review is right |
…ll paths - test_sparse_backend_prefill_correctness gains the GLM-5.3-Flash NoPE layout, so dense MHA, masked MHA and masked MHA with chunked context are checked against the SDPA reference with qk_rope_head_dim == 0. - _is_masked_mha_available: allow-list test over the three supported geometries plus quantized-KV, FA3 and mismatched-dimension negatives. - FlashAttnPrefillBackend.validate_configuration accepts (256, 0, 256) and still rejects (256, 64, 256). - _concat_k_nope_k_pe equals torch.cat of k_nope and the broadcast k_pe, and returns k_nope itself (no allocation) when there is no RoPE part. Signed-off-by: Jared Wen <jaredwen@inferact.ai> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Member
|
/ci run |
ZJY0516
enabled auto-merge (squash)
September 9, 2026 07:10
|
✅ Triggered Buildkite CI #87857 for commit |
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
…sked-mha # Conflicts: # tests/v1/attention/test_sparse_mla_backends.py
auto-merge was automatically disabled
September 14, 2026 08:17
Head branch was pushed to by a user without write access
Contributor
Author
|
/ci run |
|
✅ Triggered Buildkite CI #88766 for commit |
ZJY0516
enabled auto-merge (squash)
September 14, 2026 08:47
Shreya-gaur
pushed a commit
to Shreya-gaur/vllm_private
that referenced
this pull request
Sep 14, 2026
…56, 0, 256) layout + skip the NoPE K concat (vllm-project#55738) Signed-off-by: Jared Wen <jaredwen@inferact.ai> Signed-off-by: Jared Wen <w13431838023@gmail.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Purpose
One of three independent GLM-5.3-Flash perf PRs
On main, GLM-5.3-Flash's MLA layout (qk_nope 256, qk_rope 0, v 256) is unknown to the FlashAttention prefill backend and to the masked-MHA allow-list, so
MLACommonImpllogsNo MLA prefill backend supports this modeland every prefill token goes through the per-token top-k MQA kernel. Profiling shows that kernel is KV-gather bound at the HBM roofline: 5.9 ms/layer per 16k-token chunk, and still 1.6 ms/layer for 8x2048 where dense MHA needs ~0.2 ms.Three commits:
FlashAttnPrefillBackend(7 lines). Same lines as in the SM90 wiring PR [perf] wire FA and FlashMLA for sm90 GLM5Next NoPE SparseMLA #55385; kept separate so it can be dropped once that lands. This alone gives the dense-MHA prefill path for sequences ≤ index_topk (2048)._is_masked_mha_availableonly knew (128, 512, 128, 64, 128) and (64, 512, 192, 64, 256); GLM-5.3-Flash is (64, 512, 256, 0, 256). The kernels are identical for qk_head_dim 256 / v 256, so sequences between 2k and the existing FLASHINFER_MLA_SPARSE thresholds (36k at TP4, 64k at TP8) now use FA4 masked MHA instead of the per-token top-k MQA kernel._concat_k_nope_k_pereturnsk_nopedirectly when there is no RoPE part instead of allocating and copying a same-sized tensor (134 MB/layer per 16k chunk).Test Plan
0 prefix-cache hit
lm_eval --model local-completions --tasks gsm8k --num_fewshot 5 --gen_kwargs temperature=0(1319 questions)Test Result
Performance of this PR alone
Measured on top of #55736 + #55737 (main + #55736 + #55737) vs the same + this PR, 4x GB300, TP4,
--attention-backend FLASHINFER_MLA_SPARSE --max-model-len 69632 --max-num-seqs 256 --max-num-batched-tokens 16384, prefix caching disabled (hit rate 0.0% checked in the server log),vllm bench serverandom dataset with warmups, back-to-back in the same session. Decode tok/s is the steady-state window value; TPOT is the per-request median.Prefill up to the masked-MHA threshold (≤ 36k at TP4): −8.5% TTFT for 8x2048 (dense MHA), −7.4% for 8x32k (masked MHA), −4.4% TTFT in the 32k-context serving point. 2x8192 shows −4.6%, but that 16-request point has a ~4% run-to-run noise floor (last column of the ablation below), so treat it as indicative only. 64k prompts are above the threshold and stay on the sparse kernel; the −3.5% there (0.0% noise floor) is the K-concat removal alone. Decode is unchanged (±0.8%).
Accuracy (per build, same session)
Each rung of the ablation was also evaluated on its own: gsm8k (1319 questions, 5-shot, greedy, lm_eval
local-completions), prompt-logprob agreement on real 4k/12k/30k/60k prompts (mean |Δ logprob| per token vs the main build; the "main again" row is the run-to-run noise floor of this FP8 model), and a needle-in-a-haystack retrieval set (12 codes per length at 6k/16k/30k tokens, greedy,reasoning_effort=low), which exercises the 2k-36k prefill range where #55738 switches kernels.Caveat: runtime JIT
The first request in each new sequence-length bucket JIT-compiles the FA4 CuTeDSL
mask_modvariant (BlackwellFusedMultiHeadAttentionForward) plus the Triton_scatter_topk_single_req_kernel(~7 s each;jit_monitorwarns).FA4MLAPrefillKernel.get_warmup_keysdoes not cover mask_mod, so this is a pre-existing gap for every masked-MHA model; a warmup for the masked variants should be a follow-up before relying on this in production.Duplicate-work check
gh pr list --repo vllm-project/vllm --state open --search "NoPE masked MHA 256"/"GLM-5.3-Flash prefill": #55385 covers the FA dims (commit 1 here is the same 7 lines); #54951 / #55222 / #55543 touch the indexer prefill sharding, the indexer workspace and SM90 fp8 KV. None adds the masked-MHA layout or the K-concat skip.flashinfer_mla_sparse_sm120.py); no overlap with the masked-MHA allow-list or the K-concat skip.Tests
tests/v1/attention/test_sparse_mla_backends.py::test_sparse_backend_prefill_correctness[glm53_flash_nope_hd256_v256-*]: dense MHA, masked MHA and masked MHA with chunked context for the NoPE(64, 256, 0, 256)layout against the per-token SDPA reference.tests/v1/attention/test_sparse_mla_backends.py::test_is_masked_mha_available_model_dims: allow-list accepts the DeepSeek-V3.2, GLM-5 and NoPE GLM-5.3-Flash geometries and rejects quantized KV, FA3,(256, 64, 256)and a different head count.tests/v1/attention/test_mla_prefill_selector.py::TestBackendValidation::test_flash_attn_accepts_glm53_flash_nope_dimensions:FlashAttnPrefillBackend.validate_configurationaccepts(256, 0, 256)and still rejects(256, 64, 256).tests/v1/attention/test_mla_backends.py::test_concat_k_nope_k_pe_matches_torch_cat[rope|nope]:_concat_k_nope_k_peequalstorch.catofk_nopeand the broadcastk_pe; with no RoPE part it returnsk_nopeitself (samedata_ptr).tests/kernels/test_mhc_kernels.py,tests/kernels/test_glm5next_kda_recurrent_strided.pyunaffected; prompt-logprob A/B on real 4k/12k/30k/60k prompts vs main: mean|Δ| 0.06-0.12 vs a main-vs-main noise floor of 0.04-0.11 (≈2.5x at 4k where dense/masked MHA replaces the sparse kernel), identical next tokens; gsm8k 1319 questions (5-shot, greedy): main 93.03% ± 0.70, full series 92.65% ± 0.72 (measured on the three PRs together).pre-commit run ruff-check / ruff-format: passed.Notes for review
AI assistance: developed with Claude and reviewed by the submitter.
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.