[#14828][feat] AutoDeploy: support multi KV cache memory pool in trtllm attention - #14911
Merged
Merged
Conversation
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #52059 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #52059 [ run ] completed with state
|
Contributor
📝 WalkthroughWalkthroughTRT-LLM attention in AutoDeploy now supports multiple KV cache memory pools for non-uniform sliding-window models by introducing a backend-capability contract, implementing per-cache_loc buffer allocation in the kernel planner, removing uniform KV cache restrictions, and providing executor logic to use cyclic-view vs. host-sliced metadata computation based on kernel capability. ChangesMulti-pool KV cache support for TRT-LLM with cyclic-SWA contract
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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: 1
🧹 Nitpick comments (2)
tests/unittest/auto_deploy/singlegpu/transformations/library/test_kv_cache_trtllm_multipool.py (1)
119-205: 🏗️ Heavy liftExercise the real executor path and a multi-page case here.
This helper hand-builds
cache_loc_per_pool/seq_len_with_cache_per_pooland also keepstokens_per_block == max_seq_len, so both tests stay in a single-page regime. That means a regression intensorrt_llm/_torch/auto_deploy/shim/ad_executor.py::_prepare_inputs()or in the multi-page cyclic staging/block-table path would still pass here. Please drive one case through the executor staging path and shrinktokens_per_block(or growseq_len) so the SWA pool spans multiple pages. Coverage is currently insufficient for the executor change. As per coding guidelines,tests/**: Act as a QA engineer reviewing test changes and coverage for TensorRT-LLM.🤖 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/unittest/auto_deploy/singlegpu/transformations/library/test_kv_cache_trtllm_multipool.py` around lines 119 - 205, The test currently hand-constructs cache metadata and keeps tokens_per_block == max_seq_len so it never exercises the multi-page/executor staging path; change the setup in _build_and_stage to drive at least one case through the real executor staging (instead of fully hand-building cache_loc_per_pool/seq_len_with_cache_per_pool) and make the SWA pool span multiple pages by reducing tokens_per_block (or increasing seq_len) so tokens_per_block < seq_len (e.g., tokens_per_block much smaller than seq_len) which forces multi-page behavior; ensure you still call the same CachedSequenceInterface instance (cm) and invoke the graph/executor path via gm(**cm.named_args) so ad_executor::_prepare_inputs() and the cyclic staging/block-table code are exercised rather than bypassed.tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_trtllm_attention_op.py (1)
903-903: ⚡ Quick winAdd type annotations to the newly introduced test helper/methods.
The new methods in this class are currently untyped (e.g., Line 903, Line 921, Line 953). Please annotate parameters and return types (
-> Nonefor tests) to match repo typing guidance.Suggested patch
@@ `@staticmethod` - def _host_prepare(num_seq, max_seq_len, tokens_per_block, max_batch_size, device): + def _host_prepare( + num_seq: int, + max_seq_len: int, + tokens_per_block: int, + max_batch_size: int, + device: str, + ) -> torch.Tensor: @@ - def test_per_group_buffers_are_distinct_and_not_clobbered(self, device): + def test_per_group_buffers_are_distinct_and_not_clobbered(self, device: str) -> None: @@ - def test_same_cache_loc_returns_stable_buffer(self, device): + def test_same_cache_loc_returns_stable_buffer(self, device: str) -> None:As per coding guidelines: "Always annotate functions; make the return type
Noneif the function does not return anything."Also applies to: 921-970
🤖 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/unittest/auto_deploy/singlegpu/custom_ops/attention/test_trtllm_attention_op.py` at line 903, Several newly added test helper methods (including _host_prepare and the other helpers introduced around lines 921 and 953) lack type annotations; update each function signature to annotate all parameters with appropriate types and set the return type to None (e.g., def _host_prepare(num_seq: int, max_seq_len: int, tokens_per_block: int, max_batch_size: int, device: torch.device) -> None), following the repo typing conventions so the test helpers are fully typed.
🤖 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
`@tests/unittest/auto_deploy/singlegpu/transformations/library/test_kvcache_vswa_metadata.py`:
- Line 222: The test unpacks an unused variable `gm` from the call gm, info, cm
= _run_transform(backend=backend) which triggers Ruff RUF059; fix it by removing
the unnecessary binding or replace `gm` with `_` (e.g., _, info, cm =
_run_transform(...)) in the test_kvcache_vswa_metadata parametrized test so only
used values (`info`, `cm`) are kept; update any local references if they relied
on `gm` and ensure `_run_transform` call remains unchanged.
---
Nitpick comments:
In
`@tests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_trtllm_attention_op.py`:
- Line 903: Several newly added test helper methods (including _host_prepare and
the other helpers introduced around lines 921 and 953) lack type annotations;
update each function signature to annotate all parameters with appropriate types
and set the return type to None (e.g., def _host_prepare(num_seq: int,
max_seq_len: int, tokens_per_block: int, max_batch_size: int, device:
torch.device) -> None), following the repo typing conventions so the test
helpers are fully typed.
In
`@tests/unittest/auto_deploy/singlegpu/transformations/library/test_kv_cache_trtllm_multipool.py`:
- Around line 119-205: The test currently hand-constructs cache metadata and
keeps tokens_per_block == max_seq_len so it never exercises the
multi-page/executor staging path; change the setup in _build_and_stage to drive
at least one case through the real executor staging (instead of fully
hand-building cache_loc_per_pool/seq_len_with_cache_per_pool) and make the SWA
pool span multiple pages by reducing tokens_per_block (or increasing seq_len) so
tokens_per_block < seq_len (e.g., tokens_per_block much smaller than seq_len)
which forces multi-page behavior; ensure you still call the same
CachedSequenceInterface instance (cm) and invoke the graph/executor path via
gm(**cm.named_args) so ad_executor::_prepare_inputs() and the cyclic
staging/block-table code are exercised rather than bypassed.
🪄 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: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 9d508c5e-4fd7-4873-8e83-2745e29257e5
📒 Files selected for processing (10)
tensorrt_llm/_torch/auto_deploy/custom_ops/attention/trtllm_attention.pytensorrt_llm/_torch/auto_deploy/custom_ops/attention_interface.pytensorrt_llm/_torch/auto_deploy/llm_args.pytensorrt_llm/_torch/auto_deploy/shim/ad_executor.pytensorrt_llm/_torch/auto_deploy/shim/interface.pytensorrt_llm/_torch/auto_deploy/transform/library/kvcache.pytests/unittest/auto_deploy/singlegpu/custom_ops/attention/test_trtllm_attention_op.pytests/unittest/auto_deploy/singlegpu/shim/test_ad_executor_swa_eviction.pytests/unittest/auto_deploy/singlegpu/transformations/library/test_kv_cache_trtllm_multipool.pytests/unittest/auto_deploy/singlegpu/transformations/library/test_kvcache_vswa_metadata.py
Collaborator
Author
|
/bot run |
…n trtllm attention Enable multiple KV cache memory pools in AutoDeploy's trtllm attention so non-uniform sliding-window models (e.g. gpt-oss) work on the trtllm backend. The window-group machinery already existed for triton/flashinfer; this removes the three trtllm-specific blockers: - Gate: requires_uniform_kv_caches now returns False, so the unified KVCacheManager may host more than one pool for trtllm. - Per-group block_offsets: the trtllm planner keeps an address-stable block_offsets buffer per KV window group (keyed by the group's cache_loc input pointer) so per-group prepare_trtllm_metadata invocations no longer clobber a single shared buffer. - Cyclic-window staging: the trtllm kernel masks the sliding window internally via cyclic indexing, so the executor passes the full per-window block table and global KV length (mirroring the PyTorch backend) instead of host-slicing. A new AttentionDescriptor.kernel_handles_cyclic_swa() capability (True only for trtllm) is plumbed through the kvcache transform and CachedSequenceInterface. Adds unit coverage for per-group buffers, cyclic-view staging, the gate, the backend plumbing, and a forward-level two-pool trtllm test. Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com>
…st for multi KV pool trtllm attention now supports multiple KV cache memory pools, so requires_uniform_kv_caches defaults to False for all backends. Update the test assertion to match the new contract. Signed-off-by: egeva <19514940+MrGeva@users.noreply.github.com>
MrGeva
force-pushed
the
eg/ad-trtllm-multipool-kv
branch
from
June 7, 2026 09:01
d87bc68 to
da0dcfd
Compare
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #52566 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #52567 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #52566 [ run ] completed with state |
Collaborator
|
PR_Github #52567 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #52582 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #52582 [ run ] completed with state
|
…SWA test The triton_paged backend was consolidated into the 'triton' backend (registered in triton_attention.py), so the parametrized test_vswa_sets_kernel_handles_cyclic_swa case failed in CI with 'Attention source triton_paged not registered'. Rename the case to 'triton', matching the sibling tests in this file and the current attention registry. The triton descriptor inherits the base kernel_handles_cyclic_swa() default (False), so expected_cyclic stays False. Signed-off-by: egeva <19514940+MrGeva@users.noreply.github.com>
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #52590 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #52590 [ run ] completed with state
|
taylor-yb-lee
approved these changes
Jun 7, 2026
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #52685 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #52685 [ run ] completed with state
|
…navailable In the standalone auto_deploy CI environment, TRT-LLM native bindings are absent so the 'trtllm' attention backend never registers. Three tests then fail with 'Attention source trtllm not registered': - test_vswa_sets_kernel_handles_cyclic_swa[trtllm-True]: add a pytest.skip() guard when trtllm_ops_available() is False; the [triton-False] case continues to run standalone. - test_kv_cache_trtllm_multipool.py (both tests): extend the existing pytestmark skipif to also require trtllm_ops_available(). - create_standalone_package.py: exclude test_kv_cache_trtllm_multipool.py from the standalone package (same treatment as test_gemm_fusion_trtllm.py and test_kv_cache.py which have the same dependency). Signed-off-by: Eran Geva <19514940+MrGeva@users.noreply.github.com>
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #52761 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #52761 [ run ] completed with state |
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
3 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.
Summary
Enables multiple KV cache memory pools in AutoDeploy's
trtllmattention backend, so non-uniform sliding-window models (e.g. gpt-oss, which alternates sliding-window and full-attention layers) run on the trtllm backend. Fixes #14828.The AutoDeploy multi-pool ("window group") machinery already existed for the
triton/flashinferbackends. This PR removes the three trtllm-specific blockers:llm_args.py):requires_uniform_kv_cachesnow returnsFalse, so_identify_managed_kv_resourcesno longer raises when more than one pool is present. AutoDeploy's trtllm path uses per-layer KV cache views, so theblock_offset_multiplieris uniformlykv_factor=2regardless of pool count.block_offsets(custom_ops/attention/trtllm_attention.py): the planner previously wrote a single module-levelblock_offsetsbuffer, which the kvcache transform's per-groupprepare_trtllm_metadatainvocations clobbered. It now keeps one address-stable buffer per KV window group, keyed by the group'scache_locinput pointer (same pattern as the existing per-layer cache); group 0 reuses the pre-allocatedreset()buffer. CUDA-graph safe (lazy allocation gated to warm-up).shim/ad_executor.py): the trtllm kernel applies the sliding-window mask internally via cyclic KV indexing, so it needs the full per-window block table + global KV length — like the PyTorch backend — rather than the host-sliced window view used by triton/flashinfer. A newAttentionDescriptor.kernel_handles_cyclic_swa()capability (True only for trtllm) is plumbed through the kvcache transform →CachedSequenceInterface→ executor, which then uses the new_compute_cyclic_full_viewhelper. triton/flashinfer keep host-slicing unchanged.The kvcache transform needed only a single setter call — it was already backend-agnostic for grouping.
Test coverage
test_trtllm_attention_op.py: per-groupblock_offsetsbuffers are distinct and not clobbered across groups; samecache_locreturns an address-stable buffer.test_kvcache_vswa_metadata.py: the transform recordskernel_handles_cyclic_swaper backend (True for trtllm, False for triton).test_ad_executor_swa_eviction.py:_compute_cyclic_full_viewpasses the full block table + global length; cyclic vs window-local staging diverge under eviction as expected; the multi-pool gate is allowed by default and the uniform-cache mechanism is still enforced when explicitly requested.test_kv_cache_trtllm_multipool.py: forward-level test — a real two-distinct-window model with the trtllm cached op matches the eager reference exactly (validates pool routing + per-group buffers + cyclic staging), plus a smoke check for the SWA-engaged path.No regressions in
test_kv_cache.py(27 passed) or the full trtllm op suite (32 passed). Allpre-commithooks pass.Validation
RuntimeError: KV resources are not uniformbefore this change.Performance Impact
★ MEASURED THROUGHPUT (live A/B, node gpu-120, 1×B200, trtllm-bench _autodeploy)
A/B = identical config, only
AD_FORCE_SINGLE_KV_POOLtoggled (0=multi-pool/SWA on,1=single-pool/SWA off). 512 reqs @ concurrency 256. Both modes got the SAME KV
budget (27.96 GB here) so it is apples-to-apples.
ISL 16384 / OSL 1024:
Summary by CodeRabbit
New Features
Tests