Repository navigation
perf(attention): reuse TRT-LLM-gen counter buffers - #3949
Conversation
Preallocate the direct GQA and MLA benchmark counter sidecars outside timed calls. Add caller-owned MLA buffer reuse with validation and regression coverage for repeated launches.
|
No actionable comments were generated in the recent review. ð âđïļ Recent review infoâïļ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: ð Files selected for processing (5)
ð§ Files skipped from review as they are similar to previous changes (2)
ð WalkthroughWalkthroughThe TRT-LLM decode and prefill paths now resolve reusable multi-CTA KV counter buffers through shared validation logic. MLA dispatch, benchmarks, autotuning, and tests support caller-provided buffers, repeated execution, output reuse, and invalid-buffer checks. ChangesTRT-LLM multi-CTA counter buffer
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant trtllm_batch_decode_with_kv_cache_mla
participant _resolve_trtllm_gen_multi_ctas_kv_counter_buffer
participant TrtllmGenMlaDecodeRunner
participant GPU
Caller->>trtllm_batch_decode_with_kv_cache_mla: provide optional counter buffer
trtllm_batch_decode_with_kv_cache_mla->_resolve_trtllm_gen_multi_ctas_kv_counter_buffer: validate or allocate buffer
_resolve_trtllm_gen_multi_ctas_kv_counter_buffer-->>trtllm_batch_decode_with_kv_cache_mla: return resolved buffer
trtllm_batch_decode_with_kv_cache_mla->>TrtllmGenMlaDecodeRunner: pass buffer for TRT-LLM execution
TrtllmGenMlaDecodeRunner->>GPU: execute decode and reuse counter buffer
GPU-->>Caller: return output
Possibly related PRs
Suggested labels: Suggested reviewers: ðĨ Pre-merge checks | â 4 | â 1â Failed checks (1 warning)
â Passed checks (4 passed)
âĻ Finishing Touchesð Generate docstrings
𧊠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 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces support for a caller-owned, reusable multi_ctas_kv_counter_buffer in the trtllm-gen backend for both GQA and MLA decoding. It adds validation logic to ensure the buffer is contiguous, on the correct device, and of sufficient size, and integrates it into the decoding runners and API entry points. Benchmarks and unit tests have been updated to verify correct behavior and error handling. No review comments were provided, so there is no additional feedback.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
ð§đ Nitpick comments (1)
flashinfer/mla/_core.py (1)
2209-2235: ð Maintainability & Code Quality | ðĩ Trivial | ðĪ Low valueConsider placing the resolver in
flashinfer/utils.pyfor cross-module reuse.The validation logic here (device match, contiguity, minimum byte size) duplicates the inline validation in
flashinfer/prefill.py:251-274. Placing this function alongsideget_trtllm_gen_multi_ctas_kv_counter_bytesand_get_trtllm_gen_multi_ctas_kv_counter_bufferinflashinfer/utils.pywould let both paths share one implementation and keep the added contiguity check consistent everywhere.ðĪ 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 2209 - 2235, Move _resolve_trtllm_gen_multi_ctas_kv_counter_buffer into flashinfer/utils.py alongside get_trtllm_gen_multi_ctas_kv_counter_bytes and _get_trtllm_gen_multi_ctas_kv_counter_buffer, then update both the _core.py and prefill.py call paths to reuse it. Preserve the existing device, contiguity, and minimum-byte validations consistently across both paths.
ðĪ 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 `@flashinfer/mla/_core.py`:
- Around line 3136-3156: The MLA decode autotune path must not reuse a
caller-owned multi_ctas_kv_counter_buffer when profiling batch sizes larger than
query.size(0). Update the flow around
_resolve_trtllm_gen_multi_ctas_kv_counter_buffer and TrtllmGenMlaDecodeRunner to
fall back to an internally allocated buffer during autotune, or explicitly
reject/document the combination with flashinfer.autotune(True).
---
Nitpick comments:
In `@flashinfer/mla/_core.py`:
- Around line 2209-2235: Move _resolve_trtllm_gen_multi_ctas_kv_counter_buffer
into flashinfer/utils.py alongside get_trtllm_gen_multi_ctas_kv_counter_bytes
and _get_trtllm_gen_multi_ctas_kv_counter_buffer, then update both the _core.py
and prefill.py call paths to reuse it. Preserve the existing device, contiguity,
and minimum-byte validations consistently across both paths.
ðŠ 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
Run ID: 3e78cbb8-9e03-4a59-9ff5-a663d845e82f
ð Files selected for processing (4)
benchmarks/routines/attention.pyflashinfer/mla/_core.pytests/attention/test_trtllm_gen_attention_decode.pytests/attention/test_trtllm_gen_mla.py
Use runner-owned counter storage for synthetic autotune profiles, then pass a caller-owned buffer only to the selected TRT-LLM runner's final request. Add focused coverage for profiling reuse and final-call buffer identity.
There was a problem hiding this comment.
ð§đ Nitpick comments (1)
tests/attention/test_trtllm_gen_mla.py (1)
1205-1249: ð Maintainability & Code Quality | ðĩ Trivial | ⥠Quick winDecouple the buffer-validation assertions from the parametrized matrix
TheseValueErrorchecks only run for(batch_size=1, q_len_per_request=1), so coverage silently disappears if that combination is removed from the parametrization later. Move them into a dedicated non-parametrized test or make the sentinel case explicit.ðĪ 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 `@tests/attention/test_trtllm_gen_mla.py` around lines 1205 - 1249, Move the buffer-validation checks out of the batch-size/query-length parametrized test and into a dedicated non-parametrized test, reusing the necessary setup for query, KV cache, workspace, and counter buffers. Ensure the new test independently covers both the undersized and non-contiguous multi_ctas_kv_counter_buffer cases without relying on the (1, 1) sentinel combination.
ðĪ 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.
Nitpick comments:
In `@tests/attention/test_trtllm_gen_mla.py`:
- Around line 1205-1249: Move the buffer-validation checks out of the
batch-size/query-length parametrized test and into a dedicated non-parametrized
test, reusing the necessary setup for query, KV cache, workspace, and counter
buffers. Ensure the new test independently covers both the undersized and
non-contiguous multi_ctas_kv_counter_buffer cases without relying on the (1, 1)
sentinel combination.
âđïļ Review info
âïļ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 6a78f34d-c9ee-43f3-99f6-215057d58f18
ð Files selected for processing (3)
tests/attention/test_trtllm_gen_attention_decode.pytests/attention/test_trtllm_gen_mla.pytests/autotuner/test_autotuner_mla_decode.py
ðĪ Files with no reviewable changes (1)
- tests/attention/test_trtllm_gen_attention_decode.py
|
Verified the autotune ownership contract at head
|
|
/bot run tests/attention tests/autotuner/test_autotuner_mla_decode.py |
|
[FAILED] Pipeline #58004181: 14/20 passed |
ð Description
Fix the B300 TRT-LLM-native decode regression caused by allocating and zeroing the multi-CTA KV counter sidecar inside every timed call.
Fresh validation at PR head
78b622dareran the exact B300/SM103 GPT-OSS case against the pre-regression parent using one warm process and five measured processes per ref. Median TRT-LLM-native throughput was15.330 TFLOP/sfor this PR versus15.313 TFLOP/sfor the parent. Reconstructed from the fixed operation count, latency was0.008755233 msversus0.008764953 ms, making this PR0.110894%faster. All 12 benchmark processes passed refcheck; both refs used the same CUDA-event fallback because CUPTI was unavailable.ð Related Issues
ð 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.).Validation includes the focused B300 correctness suite (
28 passed), the broader decode/MLA/prefill regression matrix (1392 passed), focused B200 caller-buffer validation (6 passed), targeted B200 autotune/final-call validation (7 passed), and the fresh exact B300 parent-versus-current benchmark above. Repository and changed-file pre-commit coverage passed, including mypy and Ruff.Reviewer Notes
Please focus on the caller-owned buffer lifetime and autotune ownership contract. Synthetic autotune profiles use reusable runner-owned storage so profile batch buckets can exceed the real request batch safely. Only the selected TRT-LLM runner's final request receives caller-owned storage. Caller buffers must remain alive for all launches or graph replays and must not be shared by overlapping streams or graphs.
Summary by CodeRabbit