[TRTLLM-12339][feat] add dual KV cache resource management - #15400
[TRTLLM-12339][feat] add dual KV cache resource management#15400cascade812 wants to merge 1 commit into
Conversation
|
/bot run |
|
PR_Github #54497 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughThis PR adds a dedicated cross KV cache pool for encoder-decoder models. It introduces ChangesDual-pool (self + cross) KV cache for encoder-decoder models
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 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}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unittest/_torch/executor/test_dual_pool_kv_cache.py (1)
493-512: ⚡ Quick winAssert cross-pool registration in
resourcesfor both non-estimation and estimation paths.These tests currently prove the constructor call but not the actual resource-map wiring contract. Add an assertion that
resources[ResourceManagerType.CROSS_KV_CACHE_MANAGER]is the created cross manager to catch silent registration regressions.Suggested test tightening
- creator._create_cross_kv_cache_manager = Mock(return_value=Mock()) + cross_mgr = Mock() + creator._create_cross_kv_cache_manager = Mock(return_value=cross_mgr) resources = {} creator.build_managers(resources, estimating_kv_cache=False) creator._create_cross_kv_cache_manager.assert_called_once() + assert resources[ResourceManagerType.CROSS_KV_CACHE_MANAGER] is cross_mgr- creator._create_cross_kv_cache_manager = Mock(return_value=Mock()) + cross_mgr = Mock() + creator._create_cross_kv_cache_manager = Mock(return_value=cross_mgr) resources = {} creator.build_managers(resources, estimating_kv_cache=True) creator._create_cross_kv_cache_manager.assert_called_once() + assert resources[ResourceManagerType.CROSS_KV_CACHE_MANAGER] is cross_mgrAs per coding guidelines, tests under
tests/**should include actionable coverage checks and clearly verify wiring behavior.Also applies to: 513-537
🤖 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/_torch/executor/test_dual_pool_kv_cache.py` around lines 493 - 512, The test method test_build_managers_registers_cross_pool_for_enc_dec currently verifies that _create_cross_kv_cache_manager is called, but does not verify that the created manager is actually registered in the resources dictionary with the correct key. After calling creator.build_managers(resources, estimating_kv_cache=False), add an assertion that confirms resources[ResourceManagerType.CROSS_KV_CACHE_MANAGER] is set to the mock manager object that was returned by _create_cross_kv_cache_manager (capture the mock's return value and compare it to the resources entry). This same verification pattern needs to be applied to the sibling test at lines 513-537 as well.Source: Coding guidelines
🤖 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/unittest/_torch/executor/test_dual_pool_kv_cache.py`:
- Around line 493-512: The test method
test_build_managers_registers_cross_pool_for_enc_dec currently verifies that
_create_cross_kv_cache_manager is called, but does not verify that the created
manager is actually registered in the resources dictionary with the correct key.
After calling creator.build_managers(resources, estimating_kv_cache=False), add
an assertion that confirms resources[ResourceManagerType.CROSS_KV_CACHE_MANAGER]
is set to the mock manager object that was returned by
_create_cross_kv_cache_manager (capture the mock's return value and compare it
to the resources entry). This same verification pattern needs to be applied to
the sibling test at lines 513-537 as well.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: c1f0ee35-e93a-4caf-8fef-2ea2105b0275
📒 Files selected for processing (4)
tensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/resource_manager.pytests/unittest/_torch/executor/test_dual_pool_kv_cache.pytests/unittest/_torch/executor/test_kv_cache_budget_split.py
|
PR_Github #54497 [ run ] completed with state
|
|
/bot run |
|
PR_Github #54563 [ run ] triggered by Bot. Commit: |
|
PR_Github #54563 [ run ] completed with state
|
|
/bot run |
|
PR_Github #54585 [ run ] triggered by Bot. Commit: |
|
PR_Github #54585 [ run ] completed with state
|
|
/bot run |
|
PR_Github #54598 [ run ] triggered by Bot. Commit: |
|
PR_Github #54598 [ run ] completed with state
|
|
/bot run |
|
PR_Github #54605 [ run ] triggered by Bot. Commit: |
|
PR_Github #54605 [ run ] completed with state
|
|
/bot run |
|
PR_Github #54612 [ run ] triggered by Bot. Commit: |
|
PR_Github #54612 [ run ] completed with state
|
|
/bot run |
|
PR_Github #54616 [ run ] triggered by Bot. Commit: |
|
PR_Github #54616 [ run ] completed with state
|
|
/bot run |
|
PR_Github #54621 [ run ] triggered by Bot. Commit: |
130e0fd to
f875cb6
Compare
|
/bot run |
|
PR_Github #54632 [ run ] triggered by Bot. Commit: |
|
PR_Github #54621 [ run ] completed with state |
|
PR_Github #54632 [ run ] completed with state
|
Signed-off-by: Guiju Zhang <guijuz@nvidia.com>
f875cb6 to
c6a5239
Compare
|
/bot run |
|
PR_Github #54661 [ run ] triggered by Bot. Commit: |
|
PR_Github #54661 [ run ] completed with state
|
|
Close this PR as #13919 has been merged into main. |
2 participants
Description
Split out the dual KV-cache resource-management slice from #13919.
This PR adds the PyTorch executor resource plumbing needed to maintain a separate cross-attention KV cache pool for encoder-decoder models:
ResourceManagerType.CROSS_KV_CACHE_MANAGERThis intentionally leaves model enablement, encoder execution, BART/T5 integration, and beam-width cross-KV block-offset expansion for follow-up split PRs.
Test Coverage
PYTHONPYCACHEPREFIX=/private/tmp/trtllm_pycache python3 -m py_compile tensorrt_llm/_torch/pyexecutor/_util.py tensorrt_llm/_torch/pyexecutor/resource_manager.py tests/unittest/_torch/executor/test_kv_cache_budget_split.py tests/unittest/_torch/executor/test_dual_pool_kv_cache.pygit diff --cached --checkNot run locally:
pytest tests/unittest/_torch/executor/test_kv_cache_budget_split.py tests/unittest/_torch/executor/test_dual_pool_kv_cache.pybecause this local system Python does not havepytestinstalled.PR Checklist
Summary by CodeRabbit
New Features
Refactor
Tests