feat: support native qk_rope_head_dim=0 sparse MLA decode in trtllm-gen - #4108
Merged
saltyminty merged 4 commits intoJul 30, 2026
Merged
Conversation
Add a native decode path for the sparse MLA shape with no rotary tail (kv_lora_rank=512, qk_rope_head_dim=0). The query carries no RoPE component and both KV TMA descriptors address a single 512-wide cache pool, so the kernel needs the per-request active lengths to bound the sparse gather. - csrc/trtllm_fmha_kernel_launcher.cu: add an optional sparse_mla_top_k_lens argument to trtllm_paged_attention_decode; when the single-pool dynamic sparse MLA shape is detected, pass the key cache as the sliding-window KV pool so the kernel reads the active lengths. The launcher already rejects combining block-sparse attention with sparse MLA, so the two paths stay mutually exclusive. - flashinfer/mla/_core.py: register nope_mla_dimensions (kv_lora_rank=512, qk_rope_head_dim=0), require sparse_mla_top_k > 0 and sparse_mla_top_k_lens for this shape, thread the autotune profiling length through the decode tuning config so different top_k values key distinct configs, and expose sparse_mla_top_k_lens on trtllm_batch_decode_with_kv_cache_mla. - flashinfer/decode.py: forward the new optional argument at the two existing kernel call sites. - flashinfer/trace/templates/attention.py: declare the optional sparse_mla_top_k_lens input on the sparse MLA decode trace template.
JustinTong0323
requested review from
aleozlx,
bkryu,
cyx-6,
saltyminty,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
July 23, 2026 09:36
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
JustinTong0323
requested review from
dhiraj113,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
qiching,
qsang-nv and
samuellees
as code owners
July 23, 2026 09:36
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (2)
📝 WalkthroughWalkthroughAdds optional per-query sparse MLA top-k lengths, supports native no-RoPE MLA dimensions, validates and autotunes the new tensor, and forwards it through Python dispatch into the TRTLLM paged-attention launcher. ChangesDynamic sparse MLA decode
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant trtllm_batch_decode_with_kv_cache_mla
participant TrtllmGenMlaDecodeRunner
participant trtllm_paged_attention_decode
participant trtllm_paged_attention_launcher
Caller->>trtllm_batch_decode_with_kv_cache_mla: pass sparse_mla_top_k_lens
trtllm_batch_decode_with_kv_cache_mla->>TrtllmGenMlaDecodeRunner: provide validated tensor
TrtllmGenMlaDecodeRunner->>trtllm_paged_attention_decode: forward optional fifth input
trtllm_paged_attention_decode->>trtllm_paged_attention_launcher: forward lengths and conditional key-cache pool
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
flashinfer/mla/_core.py (1)
2175-2224: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy liftKeep flattened sparse lengths proportional to
q_len.The sweep sets
sparse_mla_top_k_lens.shape[0]to the batch bucket, but the kernel requiresbatch_size * q_len. Forq_len > 1, autotune profiles pass a too-short tensor and fail the[sumQ]validation. Generate this tensor with a constraint/initializer sized to the synthesized query’s first two dimensions instead of sweeping it as a peer batch axis.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@flashinfer/mla/_core.py` around lines 2175 - 2224, Update the sparse-length handling in the TuningConfig construction and init_sparse_top_k_lens so sparse_mla_top_k_lens is not swept as an independent batch-axis tensor. Size its flattened output from the synthesized query’s first two dimensions (batch_size * q_len), using the appropriate constraint/initializer while retaining sparse_top_k_width values and the existing non-sparse behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@csrc/trtllm_fmha_kernel_launcher.cu`:
- Around line 434-447: Update the sparse_mla_top_k_lens validation in the
optional tensor handling block to require top_k_lens.IsContiguous() before
extracting its raw pointer. Keep the existing dtype, shape, size, and device
checks unchanged so the FFI entrypoint rejects strided tensors before the kernel
reads them as contiguous.
In `@flashinfer/mla/_core.py`:
- Around line 2802-2807: Update the nearby sparse MLA documentation to
distinguish the no-RoPE TRTLLM-GEN configuration, stating that it uses
qk_rope_head_dim=0 while preserving the existing qk_rope_head_dim=64 description
for other sparse paths.
In `@flashinfer/trace/templates/attention.py`:
- Around line 2630-2632: In the sparse trace template definitions and
dispatcher, add a ragged sparse template keyed by num_query_tokens that declares
the cum_seq_lens_q layout: 3D flattened query and 2D compact block_tables.
Update dispatch so calls with cum_seq_lens_q select this template, while other
sparse calls continue using the existing template.
---
Outside diff comments:
In `@flashinfer/mla/_core.py`:
- Around line 2175-2224: Update the sparse-length handling in the TuningConfig
construction and init_sparse_top_k_lens so sparse_mla_top_k_lens is not swept as
an independent batch-axis tensor. Size its flattened output from the synthesized
query’s first two dimensions (batch_size * q_len), using the appropriate
constraint/initializer while retaining sparse_top_k_width values and the
existing non-sparse behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 6c836125-81fa-44a3-850c-1de820865c6d
📒 Files selected for processing (4)
csrc/trtllm_fmha_kernel_launcher.cuflashinfer/decode.pyflashinfer/mla/_core.pyflashinfer/trace/templates/attention.py
Comment on lines
+2630
to
+2632
| "num_query_tokens": Var( | ||
| description="Flattened number of query tokens with active top-k lengths." | ||
| ), |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Add a ragged sparse trace template.
The dispatcher selects this template for every sparse call, including cum_seq_lens_q calls. Those use a 3D flattened query and 2D compact block_tables, while this template declares 4D and 3D tensors. Add a ragged sparse template keyed by num_query_tokens and dispatch to it when cum_seq_lens_q is present.
Also applies to: 2666-2671
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@flashinfer/trace/templates/attention.py` around lines 2630 - 2632, In the
sparse trace template definitions and dispatcher, add a ragged sparse template
keyed by num_query_tokens that declares the cum_seq_lens_q layout: 3D flattened
query and 2D compact block_tables. Update dispatch so calls with cum_seq_lens_q
select this template, while other sparse calls continue using the existing
template.
…head_dim=0 docs - Add a contiguity check in the FFI entrypoint so a strided 1D tensor is not read as contiguous by the kernel (Python normalizes today, but the exported entrypoint must preserve the contract). - Clarify that the native no-RoPE path uses qk_rope_head_dim=0 while other sparse MLA paths use 64.
Contributor
|
/bot run tests/attention |
Collaborator
Collaborator
|
[SUCCESS] Pipeline #59286008: 18/18 executed test jobs passed |
The sparse_mla_top_k_lens trace template addition declared a num_query_tokens Var without a matching init kwarg, failing tests/trace/test_template_init.py on H100. The trace/AOT schema is not needed for the runtime decode path; defer trace support (including the ragged cum_seq_lens_q layout) to a dedicated follow-up.
Contributor
|
/bot run tests/attention |
Collaborator
Collaborator
|
[FAILED] Pipeline #59596789 — 16/18 executed test jobs passed Compared with nightly #59545175. Unit Tests
✅ Pass · 🟡 Old failure · ❌ New failure · ⏱ Test timeout · Multi-GPU and Multi-Node Tests — 6/6 passed
Failure detailsTimeouts, infrastructure, or incomplete jobs
|
Reconcile the native qk_rope_head_dim=0 sparse MLA decode path with upstream's refactor (launcher now carries sparse_mla_top_k_lens; DCP added to the CuTeDSL MLA path). The autotune fifth batch-swept tensor is shared: sparse top-k lengths (trtllm-gen) and DCP causal bound (cute-dsl) are mutually exclusive, routed to disjoint runners (enable_dcp forces backend=cute-dsl).
JustinTong0323
force-pushed
the
feat/native-nope-trtllm-mla
branch
from
July 29, 2026 08:45
d77a325 to
e6e38ca
Compare
Contributor
|
/bot run tests/attention |
Collaborator
Collaborator
|
CI looks good |
saltyminty
approved these changes
Jul 29, 2026
saltyminty
left a comment
saltyminty
left a comment
Collaborator
There was a problem hiding this comment.
Left one comment but approved
| if backend == "auto": | ||
| cc = get_compute_capability(query.device) | ||
| if cc[0] == 12 and sparse_mla_top_k > 0: | ||
| backend = "sparse" |
Collaborator
There was a problem hiding this comment.
Is it intended that we don't have a auto branch for nope?
Contributor
Author
There was a problem hiding this comment.
Yes, on SM100/SM103 the call intentionally stays on the generic auto path. Since the no-RoPE shape requires sparse MLA, CuTeDSL is filtered out and trtllm-gen is the only remaining runner, so a separate branch would duplicate the existing routing.
JustinTong0323
added a commit
to JustinTong0323/sglang
that referenced
this pull request
Jul 31, 2026
…ackend Consume flashinfer's new sparse_mla_top_k_lens argument so the DSA trtllm-gen backend can serve the native no-rotary-tail MLA shape (kv_lora_rank=512, qk_rope_head_dim=0). - transform_index: add prepare_trtllm_nope_sparse_metadata, a Triton kernel that derives per-query active top-k lengths from the packed page table and points fully-empty CUDA-graph padding rows at a valid dummy token (the native H512 kernel produces NaNs for empty rows). - dsa_backend: for the qk_rope_head_dim == 0 shape, build the active lengths and pass them to trtllm_batch_decode_with_kv_cache_mla. Requires flashinfer with sparse_mla_top_k_lens support (flashinfer-ai/flashinfer#4108).
This was referenced Aug 5, 2026
5 tasks
bkryu
pushed a commit
that referenced
this pull request
Sep 8, 2026
…backends (#4947) ## Summary `trtllm_batch_decode_with_kv_cache_mla` rejects the native NoPE form (`kv_lora_rank=512`, `qk_rope_head_dim=0`) at API entry unless `sparse_mla_top_k_lens` is provided. That requirement belongs to the native no-rope trtllm-gen/cute-dsl kernels (#4108), which consume the per-token active top-k length. The SM120 sparse backend (`backend="sparse"`, the v32 / GLM53_NOPE families) bounds each row by its `-1` page-table entries and never reads `sparse_mla_top_k_lens` — so the entry-level check makes the SM120 GLM-5.3-Flash NoPE path uncallable. (#4842 hit the same wall and dropped the check wholesale; this PR keeps the guard where the consuming kernels are instead.) Move the requirement past backend resolution and apply it only when a non-`sparse` backend will run. The `sparse_mla_top_k_lens` shape/dtype validation for callers that do pass it is unchanged, as is the SM100 native-NoPE contract. ## Testing - vLLM `FLASHINFER_MLA_SPARSE_SM120` + GLM-5.3-Flash-NVFP4, TP4 on 4×RTX PRO 6000 (SM120): previously raised `Native qk_rope_head_dim=0 TRTLLM-GEN MLA requires sparse_mla_top_k_lens` during CUDA graph capture; with this change the server boots and serves (companion vLLM PR: vllm-project/vllm#55277). GSM8K strict-match 0.9325. - Existing trtllm-gen NoPE callers are unaffected: the requirement still fires for `trtllm-gen` / `cute-dsl` / `xqa` / unresolved `auto` on non-SM120 parts. Signed-off-by: Zihua Wu <zihuaw@nvidia.com> <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Improved validation for sparse attention configurations with positive top-k values and per-token top-k lengths. * Updated backend-specific handling so SM120 uses per-token sequence lengths and rejects unsupported sparse top-k length settings. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Zihua Wu <zihuaw@nvidia.com> Co-authored-by: Kimi Code <noreply@moonshot.cn>
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
4 participants
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.
📌 Description
Adds a native TRTLLM-GEN sparse MLA decode path for the shape with no rotary tail:
kv_lora_rank=512,qk_rope_head_dim=0.For this shape the query carries no RoPE component, and both KV TMA descriptors address a single 512-wide cache pool, so the kernel needs the per-request active lengths to bound the sparse gather. Today
trtllm_batch_decode_with_kv_cache_mlaonly acceptsdeepseek_mla_dimensionsandsmaller_mla_dimensions, so aqk_rope_head_dim=0request is rejected as an unsupported MLA dimension.This PR registers the new dimension set and threads an optional
sparse_mla_top_k_lenstensor down to the launcher so the shape can be served natively.Changes
csrc/trtllm_fmha_kernel_launcher.cu— add an optionalsparse_mla_top_k_lensargument totrtllm_paged_attention_decode. When the single-pool dynamic sparse MLA shape is detected (sparse_mla_top_k_lenspresent and MLA decode), pass the key cache as the sliding-window KV pool so the kernel reads the active per-token lengths. The launcher already rejects combining block-sparse attention with sparse MLA (sparse_mla_top_k <= 0check), so the two stay mutually exclusive.flashinfer/mla/_core.py— registernope_mla_dimensions(kv_lora_rank=512,qk_rope_head_dim=0); requiresparse_mla_top_k > 0and asparse_mla_top_k_lenstensor for this shape; thread the autotune profiling length through the decode tuning config so differenttop_kvalues key distinct autotune configs; exposesparse_mla_top_k_lenson the publictrtllm_batch_decode_with_kv_cache_mla.flashinfer/decode.py— forward the new optional argument at the two existing kernel call sites.flashinfer/trace/templates/attention.py— declare the optionalsparse_mla_top_k_lensinput on the sparse MLA decode trace template so the trace schema matches the kernel signature.The new argument is optional and defaults to
None, so thedeepseek_mla_dimensions/smaller_mla_dimensionsdecode paths are unchanged.sparse_mla_top_k_lens(oneint32active length per query token) is supplied by the caller.🔍 Related Issues
None.
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Static checks pass (
clang-format/ruff/mypyviapre-commit). The new path has been exercised end-to-end in a downstream serving stack that computessparse_mla_top_k_lensfrom the page table and drives this decode path. Happy to add a focused in-tree unit test for theqk_rope_head_dim=0dimension registration + argument threading — see Reviewer Notes.Reviewer Notes
sparse_mla_top_k_lensis optional and defaults toNone; all existing callers and the two established MLA dimension sets keep their current behavior.sparse_mla_top_k <= 0check), so the new single-pool path cannot be entered together with block-sparse.len(inputs)==4) and a sparse request (len(inputs)==5) resolve to distinct autotune configs rather than mis-keying.Summary by CodeRabbit
New Features
Bug Fixes