Conversation
flashinfer's trtllm_batch_decode_with_kv_cache_mla instantiates a fresh runner per call, so its internal counter-buffer cache never survives across decode steps and it torch.zeros() a new one every call by default. Own a persistent buffer in TRTLLMMLABackend and DeepseekSparseAttnBackend and pass it via multi_ctas_kv_counter_buffer= instead. DSA's trtllm sparse-indexer path flattens every token in a prefill/ extend/verify chunk into its own decode-shaped batch entry, so its batch dimension can exceed max_running_requests; grow the buffer on demand there instead of assuming a fixed bound.
b8zhong
requested review from
Fridge003,
HaiShaw,
Qiaolin-Yu,
hebiao064,
ispobock and
merrymercy
as code owners
July 17, 2026 19:28
Contributor
|
Caution The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased. |
Collaborator
Author
|
/rerun-test test_dsa_glm52_nvfp4_dp_mtp.py |
b8zhong
changed the title
Contributor
|
Results for 🚀 |
added 2 commits
July 17, 2026 23:11
b8zhong
marked this pull request as draft
July 18, 2026 00:18
b8zhong
added a commit
that referenced
this pull request
Jul 21, 2026
flashinfer's trtllm_batch_decode_with_kv_cache_mla instantiates a fresh runner per call, so its internal counter-buffer cache never survives across decode steps and it torch.zeros() a new one every call by default. Own a persistent buffer in TRTLLMMLABackend and DeepseekSparseAttnBackend and pass it via multi_ctas_kv_counter_buffer= instead. DSA's trtllm sparse-indexer path flattens every token in a prefill/ extend/verify chunk into its own decode-shaped batch entry, so its batch dimension can exceed max_running_requests; grow the buffer on demand there instead of assuming a fixed bound. Cherry-picked from #31607 (squashed from 0042fd9, 121c901, a84b529, ed135a8).
Collaborator
Author
|
Move #31927 |
vincentzed
pushed a commit
to vincentzed/sglang
that referenced
this pull request
Jul 23, 2026
…buffer for DCP The persistent multi-CTAS KV counter buffer (sgl-project#31607) is sized by head count; under DCP the trtllm decode runs with dcp-widened heads, so pre-size it by num_q_heads * dcp_size and let the grow check use the per-call head count. token_capacity in metrics and the load inquirer now reports logical tokens (physical * dcp_size), matching the allocator's accounting. Fix by @b8zhong, thanks. Co-Authored-By: Claude Fable 5 <noreply@anthropic.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
2 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.
flashinfer's trtllm_batch_decode_with_kv_cache_mla instantiates a fresh runner per call (even in cuda graph). Instead use the buffer by multi_ctas_kv_counter_buffer= instead.
For prefill, it will reshape to
(num_tokens, 1, num_heads, head_dim). So we also need to grow the workspace, in case some users set chunk prefill larger than 8k. It just grows the buffer once.note: trtllm ragged prefill deepseek doesn't have this bug. It's recently introduced in #31502 after flashinfer-ai/flashinfer#3949
Main:
This branch:
CI States
Latest PR Test (Base): ❌ Run #29622595001
Latest PR Test (Extra): ❌ Run #29622594907