You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
[TRTLLM-17505][feat] KVCM2: pool SINK_REQUIRED_ONLY lifecycles like REQUIRED ones - #21793
effectivePoolReusePolicy maps SINK_REQUIRED_ONLY to REQUIRED for GPU and cold-tier pool grouping. OPTIONAL and other policies retain separate grouping keys. GPU grouping still accounts for per-pool sparsity, and cold-tier grouping still accounts for cold-page size.
The change adds two C++ unit tests for matching lifecycle shapes and OPTIONAL separation. No test execution results are available.
QA Engineer Review
Two C++ unit tests were added:
SinkRequiredOnlyWindowsShareRequiredPools checks GPU pool grouping.
SinkRequiredOnlyWindowSharesTheRequiredColdPoolGroup checks cold-tier pool grouping and verifies that three lifecycles form two groups.
No integration test-list files changed. These C++ unit tests do not require integration test-list entries. Coverage verdict: sufficient for the changed GPU and cold-tier grouping behavior.
Per-File QA Perspective
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h: Adds effectivePoolReusePolicy. Verify that the mapping changes pool grouping without changing lifecycle identity or reuse behavior.
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp: GPU pool grouping now uses the mapped policy. Verify sharing for matching shapes and continued separation when per-pool sparsity differs.
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp: Cold-tier grouping now uses the mapped policy and cold-page size. Verify sharing for matching page sizes and separation when page sizes differ.
cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp: Adds SinkRequiredOnlyWindowsShareRequiredPools to check GPU pool sharing and OPTIONAL separation. This C++ unit test does not require an integration test-list entry.
cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp: Adds SinkRequiredOnlyWindowSharesTheRequiredColdPoolGroup to check cold-tier sharing and OPTIONAL separation. This C++ unit test does not require an integration test-list entry.
Description
The KV cache manager V2 core groups life cycles into GPU pools, and into cold-tier pools, by reuse policy, page shape and sparsity. A SINK_REQUIRED_ONLY life cycle therefore never shares pools with a REQUIRED one, even when their pages
have the same shape.
SINK_REQUIRED_ONLY only changes which pages a prefix match requires: its sink blocks are required and its sliding window is optional. What happens when a page is evicted already follows the page's own life cycle, not its pool: evicting a window page leaves the prefix reusable, and evicting a sink page drops the subtree. A SINK_REQUIRED_ONLY pool also already holds both kinds of pages.
This PR keys the pools of a SINK_REQUIRED_ONLY life cycle on REQUIRED, in both the GPU and the cold-tier grouping, so a reuse policy no longer moves pages into separate memory. OPTIONAL and PRIVATE life cycles keep their own pools.
Nothing on main uses SINK_REQUIRED_ONLY yet, so main's pool layouts do not change.
The first user will be my next update to #21619, which I'll push after #20197 merges. #20197 gives each attention layer a mixed-precision high layer that keeps the first N and the last N tokens at the runtime KV type instead of NVFP4. That high layer must sit in the same pool as its layer. When the #21619 option is on, the high layer uses SINK_REQUIRED_ONLY: its first N tokens are required for reuse and its last N are not, so a prefix stays reusable after its last N tokens were evicted. Without this PR, that policy would move the high layer into a pool of its own, which #20197 does not allow.
Test Coverage
New: cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp::SinkRequiredOnlyWindowsShareRequiredPools. A REQUIRED full-attention layer and a SINK_REQUIRED_ONLY window of the same page shape share one pool group and pool; an OPTIONAL window keeps its own.
On B200, built from source:
KVCM2 C++ unit tests: 14 binaries, 162 tests passed
tests/unittest/_torch/pyexecutor/kv_cache/test_mixed_precision_kv_cache.py and tests/unittest/_torch/attention/test_mixed_precision_backends.py: 67 passed
If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.
Any new dependencies have been scanned for license and vulnerabilities
…ones
The pool-group key includes the reuse policy, so a SINK_REQUIRED_ONLY lifecycle never
shares GPU or cold-tier pools with a REQUIRED one, even when their pages have the same
shape. That policy only changes which pages a prefix match requires: its sinks are
required and its window is optional, so its pool already holds both kinds of pages.
Key the pools of a SINK_REQUIRED_ONLY lifecycle on REQUIRED, so its reuse policy does
not move its pages to separate memory. OPTIONAL and PRIVATE keep their own pools.
Signed-off-by: tianruih <tianruih@nvidia.com>
No actionable comments were generated in the recent review.
Walkthrough
Pool grouping now treats SINK_REQUIRED_ONLY as REQUIRED for GPU and cold pools. Tests verify that these lifecycles share a group and that OPTIONAL remains in a separate group.
Changes
KV cache pool reuse
Layer / File(s)
Summary
Define the pool reuse policy cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h, cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp
effectivePoolReusePolicy maps SINK_REQUIRED_ONLY to REQUIRED and leaves other reuse policies unchanged. The storage configuration comment describes the pool-sharing rule and per-pool sparsity constraint.
Apply policy to pool grouping cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp, cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp, cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp, cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp
GPU and cold pool grouping use effectivePoolReusePolicy. Tests verify that REQUIRED and SINK_REQUIRED_ONLY lifecycles share a group, while OPTIONAL uses a separate group.
NVIDIA/TensorRT-LLM#21088: Introduces the reuse policy contract used by this change’s pool grouping.
Suggested reviewers: liji-nv
Merge Risk: Minimal · up to 3f957
No merge-blocking risk was identified in the pool-sharing change.
Security architecture risk: Low · up to 3f957
The change shares memory between two reuse policies while preserving their distinct reuse rules. No new access path or weakened access control was identified, but mixed-policy behavior during failures and recovery is not fully validated.
Retained concerns
No architecture-level concerns identified.
Security review details
Security Blast Radius
inferred — The demonstrated expansion is resource-contention coupling between compatible lifecycles within a storage-manager instance. Normalization does not merge OPTIONAL or PRIVATE into REQUIRED pools. Sharing allocation and eviction capacity does not itself establish cross-tenant data reachability.
Trust Boundaries and Controls
observed — The routed public-entrypoint change is inside storage construction. The inspected caller validates configuration before creating storage, and existing validation requires sink-required-only attention to have both a sliding window and sink tokens. The full production diff introduces no new external request handler or authority-bearing interface.
Resilience and Maintainability Implications
observed — Existing failure-containment controls remain relevant after pooling: locked pages are removed from eviction eligibility; migration waits on source and destination readiness events; exception handling releases destination slots; eviction rollback requeues surviving eligible pages; repeated migration requests skip duplicate pages and pages already at the destination. These mechanisms were inspected, but complete mixed-policy fault recovery was not demonstrated.
Pre-merge checks: 5
Check name
Status
Explanation
Title check
The title clearly identifies the ticket, change type, affected component, and main change: grouping SINK_REQUIRED_ONLY lifecycles with REQUIRED pools.
Description check
The description includes the required Description, Test Coverage, and PR Checklist sections. It explains the problem, solution, affected policies, tests, validation results, and intended use case. The…
Docstring Coverage
Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentat…
Linked Issues check
Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check
Check skipped because no linked issues were found for this pull request.
Full details: Description check
Explanation
The description includes the required Description, Test Coverage, and PR Checklist sections. It explains the problem, solution, affected policies, tests, validation results, and intended use case. The checklist is included and marked as reviewed.
Full details: Docstring Coverage
Explanation
Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 3 functions across 5 files.
Finishing Touches
Generate unit tests (beta)
Create a new PR
Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
The reason will be displayed to describe this comment to others. Learn more.
I traced the merged pool groups through eviction and subtree removal, slot floors and pool ratios, the resume utilization gate, the default cold-page codec's batching, and the Python consumers of pool-group indices and stats, and found no runtime problem. poolReusePolicy is the identity for REQUIRED/OPTIONAL/PRIVATE, so existing layouts do not change. The new gtest ran and passed in A30-CPP-1 (L0 #68634), and the single-GPU stages of the latest run (L0 #68662) passed; the multi-GPU stages are only waiting for the ci: full pre-merge approved label. Counts: 0 blockers, 0 possibly blocking, 2 required, 3 nice to have.
[Required] Please link the change that actually needs this, or update the description. The description says this lets the #20197 cold-page window use SINK_REQUIRED_ONLY (#21619), and quotes a 3% -> 92% hit-rate run on a stack of the two. In both PRs as posted, though, the cold-page window never takes SINK_REQUIRED_ONLY. #20197 still builds the window's high layers with reuse_policy=ReusePolicy.REQUIRED in _add_mixed_precision_layers and requires them to sit in the source layer's pool. #21619 reads mixed_precision_reuse_last_n_from_low (default False, documented as "Mixed-precision KV only") only from mixed_precision_kv, the NVFP4 mixed-dtype path, so the cold-page window path, which has no mixed_precision_kv, stays REQUIRED. If that wiring is still local, consider landing this together with it so its end-to-end test covers the shared pool.
I traced the merged pool groups through eviction and subtree removal, slot floors and pool ratios, the resume utilization gate, the default cold-page codec's batching, and the Python consumers of pool-group indices and stats, and found no runtime problem. poolReusePolicy is the identity for REQUIRED/OPTIONAL/PRIVATE, so existing layouts do not change. The new gtest ran and passed in A30-CPP-1 (L0 #68634), and the single-GPU stages of the latest run (L0 #68662) passed; the multi-GPU stages are only waiting for the ci: full pre-merge approved label. Counts: 0 blockers, 0 possibly blocking, 2 required, 3 nice to have.
You're right, neither PR uses it yet. The change that makes the mixed-precision high layer use SINK_REQUIRED_ONLY is in my next update to #21619, along with a test. I'll push it after #20197 merges.
…Policy
Add a cold-page test that checks a SINK_REQUIRED_ONLY window shares the
REQUIRED layer's cold pool group while an OPTIONAL window keeps its own.
In the GPU grouping test, assert the pool-group count instead of a pool
index that every single-buffer layer shares. Rewrite the poolReusePolicy
comment to say why SINK_REQUIRED_ONLY shares the REQUIRED pools and
OPTIONAL does not, and point the grouping comments to it.
Signed-off-by: tianruih <tianruih@nvidia.com>
StorageManager::getBufferAttr is private, so the new cold-tier grouping
test did not compile. Read the groups with
KvCacheManager::getLifeCyclePoolGroupIndices instead; life cycles are
registered in layer order.
Signed-off-by: tianruih <tianruih@nvidia.com>
The name now shows that the policy is mapped: SINK_REQUIRED_ONLY pools
like REQUIRED. "Pool" stays in the name because
BufferConfig::effectiveReusePolicy already means the policy that reuse
follows; this function only picks the pools a lifecycle shares.
Signed-off-by: tianruih <tianruih@nvidia.com>
Automatically added "ci: full pre-merge approved" because this PR has satisfied the required GitHub review approvals. Unresolved review conversations and other required checks remain independent merge requirements.
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
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.
sunnyqgg
left a comment
Dev Engineer Review
effectivePoolReusePolicymapsSINK_REQUIRED_ONLYtoREQUIREDfor GPU and cold-tier pool grouping.OPTIONALand other policies retain separate grouping keys. GPU grouping still accounts for per-pool sparsity, and cold-tier grouping still accounts for cold-page size.The change adds two C++ unit tests for matching lifecycle shapes and
OPTIONALseparation. No test execution results are available.QA Engineer Review
Two C++ unit tests were added:
SinkRequiredOnlyWindowsShareRequiredPoolschecks GPU pool grouping.SinkRequiredOnlyWindowSharesTheRequiredColdPoolGroupchecks cold-tier pool grouping and verifies that three lifecycles form two groups.No integration test-list files changed. These C++ unit tests do not require integration test-list entries. Coverage verdict: sufficient for the changed GPU and cold-tier grouping behavior.
Per-File QA Perspective
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/lifeCycleRegistry.h: AddseffectivePoolReusePolicy. Verify that the mapping changes pool grouping without changing lifecycle identity or reuse behavior.cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storage/config.cpp: GPU pool grouping now uses the mapped policy. Verify sharing for matching shapes and continued separation when per-pool sparsity differs.cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp: Cold-tier grouping now uses the mapped policy and cold-page size. Verify sharing for matching page sizes and separation when page sizes differ.cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp: AddsSinkRequiredOnlyWindowsShareRequiredPoolsto check GPU pool sharing andOPTIONALseparation. This C++ unit test does not require an integration test-list entry.cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ColdPageTest.cpp: AddsSinkRequiredOnlyWindowSharesTheRequiredColdPoolGroupto check cold-tier sharing andOPTIONALseparation. This C++ unit test does not require an integration test-list entry.Description
The KV cache manager V2 core groups life cycles into GPU pools, and into cold-tier pools, by reuse policy, page shape and sparsity. A
SINK_REQUIRED_ONLYlife cycle therefore never shares pools with aREQUIREDone, even when their pageshave the same shape.
SINK_REQUIRED_ONLYonly changes which pages a prefix match requires: its sink blocks are required and its sliding window is optional. What happens when a page is evicted already follows the page's own life cycle, not its pool: evicting a window page leaves the prefix reusable, and evicting a sink page drops the subtree. ASINK_REQUIRED_ONLYpool also already holds both kinds of pages.This PR keys the pools of a
SINK_REQUIRED_ONLYlife cycle onREQUIRED, in both the GPU and the cold-tier grouping, so a reuse policy no longer moves pages into separate memory.OPTIONALandPRIVATElife cycles keep their own pools.Nothing on main uses
SINK_REQUIRED_ONLYyet, so main's pool layouts do not change.The first user will be my next update to #21619, which I'll push after #20197 merges. #20197 gives each attention layer a mixed-precision high layer that keeps the first N and the last N tokens at the runtime KV type instead of NVFP4. That high layer must sit in the same pool as its layer. When the #21619 option is on, the high layer uses
SINK_REQUIRED_ONLY: its first N tokens are required for reuse and its last N are not, so a prefix stays reusable after its last N tokens were evicted. Without this PR, that policy would move the high layer into a pool of its own, which #20197 does not allow.Test Coverage
cpp/tests/unit_tests/batch_manager/kvCacheManagerV2ConfigTest.cpp::SinkRequiredOnlyWindowsShareRequiredPools. AREQUIREDfull-attention layer and aSINK_REQUIRED_ONLYwindow of the same page shape share one pool group and pool; anOPTIONALwindow keeps its own.tests/unittest/kv_cache_manager_v2_tests: 595 passed, 14 skippedtests/unittest/_torch/pyexecutor/kv_cache/test_mixed_precision_kv_cache.pyandtests/unittest/_torch/attention/test_mixed_precision_backends.py: 67 passedPR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.