[None][perf] enable zero-copy token passing in KVCacheManagerV2 - #17308
Open
lowsfer wants to merge 10 commits into
Open
[None][perf] enable zero-copy token passing in KVCacheManagerV2#17308lowsfer wants to merge 10 commits into
lowsfer wants to merge 10 commits into
Conversation
lowsfer
changed the title
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughKVCacheManagerV2 adds pooled digest tokens, non-owning token spans, lazy cache-key generation, zero-copy C++ token ingestion, text-only controls, native introspection, lifecycle updates, and expanded C++ and Python validation. ChangesKVCacheManagerV2 token and reuse flow
Estimated code review effort: 4 (Complex) | ~60 minutes Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ 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: 14
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp (1)
586-599: 🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy liftRelease pages during subtree removal.
removeSubtreeandremove_subtreedetach blocks without callingreleasePages()or_release_pages(). External references can keep detached blocks alive until afterStorageManager.destroy(), leaving page-manager back-pointers and eviction state invalid. Release every block before detachment, and update the stale lifecycle comments. Fix"parent.s next map".🤖 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 `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp` around lines 586 - 599, Update removeSubtree in cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp: release each block with releasePages() before detaching it, revise the lifecycle comments, and correct “parent.s next map” to “parent’s next map”. Apply the equivalent change in remove_subtree at cpp/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py: call _release_pages() before each detachment and update its stale lifecycle comments.
🧹 Nitpick comments (8)
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py (1)
796-850: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd invalid ndarray input coverage.
This test only passes a valid C-contiguous
int32array. Add cases for a non-contiguousint32view and an unsupported dtype. These cases verify the stated binding validation before it reinterprets the buffer asTokenIdExt.Test coverage summary: insufficient.
test_int32_ndarray_ingest_matches_listcovers valid list/ndarray compatibility.TestBlockKeyHashing.test_update_int_block_matches_referencecovers the updated four-byte encoding. The validation failure paths are not covered. The supplied files do not show registration intests/integration/test_lists/test-db/ortests/integration/test_lists/qa/, so that status needs follow-up. As per path instructions, changed test functions must be checked against the appropriatetest-db/andqa/list files.🤖 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/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py` around lines 796 - 850, Extend test_int32_ndarray_ingest_matches_list with validation cases for a non-contiguous int32 ndarray view and an unsupported dtype, asserting both are rejected by the ndarray ingest/probe binding before buffer reinterpretation. Keep the existing valid contiguous int32 compatibility checks unchanged, and verify the modified test is registered in the appropriate test-db/qa list files.Source: Path instructions
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp (1)
285-297: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDrop the redundant
staticonisPrefix.
isPrefixis already inside the anonymous namespace opened at Line 282, sostaticadds nothing.♻️ Proposed change
-static bool isPrefix(TokenIdExt const* prefix, size_t prefixLen, TokenIdExt const* full, size_t fullLen) +bool isPrefix(TokenIdExt const* prefix, size_t prefixLen, TokenIdExt const* full, size_t fullLen)As per coding guidelines: "prefer anonymous namespaces over
staticfor internal-linkage functions".🤖 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 `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp` around lines 285 - 297, Remove the redundant static qualifier from the isPrefix function declaration, leaving its implementation and anonymous-namespace scope unchanged.Source: Coding guidelines
tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py (1)
68-68: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMove
Noneto the end of the union.Ruff reports RUF036 on this signature.
♻️ Proposed change
- def __init__(self, data: int | bytes | None | Sequence[int | bytes] = None) -> None: + def __init__(self, data: int | bytes | Sequence[int | bytes] | None = None) -> None:🤖 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 `@tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py` at line 68, Update the __init__ data parameter annotation so None appears at the end of the union, preserving the existing int, bytes, and Sequence[int | bytes] alternatives while resolving Ruff RUF036.Source: Linters/SAST tools
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py (1)
54-85: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winTest coverage summary (
tests/**).
- Changed test functions: none. This file only reorganizes imports.
test_reuse_scope_seeds_distinct_keysandtest_blockchain_keys_are_seeded_by_reuse_scopeare unchanged; they now resolvesequence_to_blockchain_keysfrom_block_radix_treeinstead of the package root.- Test-list registration: no new test function was added, so no entry is required in
tests/integration/test_lists/test-db/ortests/integration/test_lists/qa/. Confirm the existing entry forkv_cache_manager_v2_testsstill selects this module.- Verdict: needs follow-up.
Gaps for the changes in this cohort:
- No test asserts that the Python root key equals the C++
RootBlock::makeKeyfor the sameReuseScope. Reuse-scope salting is the tenant-isolation boundary, and this PR moved the derivation intoreuse_scope_to_bytes.test_int32_ndarray_ingest_matches_listcovers token parity only.- No test covers a token id at or above
2**31on either the Python bulk-hash path or the C++ ingest path.- No test drives a digest-bearing sequence through the
knownNoDigest/text_onlypath to confirm it is rejected rather than mis-hashed.Based on learnings: run salting tests after hashing or reuse-scope changes, event-manager tests after event changes, and both backends for shared C++/Python changes. This cohort changes hashing, reuse scope, and events, so run all three suites against both
TLLM_KV_CACHE_MANAGER_V2_BACKEND=cppandpython.As per coding guidelines: "Run unit tests with
pytest tests/unittest/for relevant changes."🤖 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/kv_cache_manager_v2_tests/test_kv_cache_salting.py` around lines 54 - 85, The salting tests do not cover cross-language root-key parity, large token IDs, or rejection of digest-bearing sequences in the knownNoDigest/text_only path. Extend the relevant tests around test_reuse_scope_seeds_distinct_keys and test_blockchain_keys_are_seeded_by_reuse_scope to compare Python roots with C++ RootBlock.makeKey for identical ReuseScope values, exercise token IDs at or above 2**31 through both backends, and verify digest-bearing input is rejected by the knownNoDigest/text_only flow. Run the relevant unit suites with both TLLM_KV_CACHE_MANAGER_V2_BACKEND=cpp and python.Sources: Coding guidelines, Path instructions, Learnings
cpp/tests/unit_tests/batch_manager/kvCacheManagerV2DigestPoolTest.cpp (1)
94-140: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for move semantics and the moved-from sentinel.
The tests cover construction, copy-clone, destruction, and reuse. They do not cover
TokenIdExt(TokenIdExt&&)oroperator=(TokenIdExt&&), which are new RAII surface. A move must transfer the slot without changingdigestPoolLiveCount(), and the moved-from handle must free nothing on destruction. Move-assignment onto a live digest handle must also release the target's slot exactly once.💚 Proposed additional test
+// Moving transfers the pool slot: no clone, no double free, and the moved-from +// handle owns nothing. +TEST(DigestPoolTest, MoveTransfersSlotWithoutCloning) +{ + size_t const baseline = detail::digestPoolLiveCount(); + Digest const bytes = makeDigest(std::byte{0x33}); + Digest const other = makeDigest(std::byte{0x77}); + { + TokenIdExt source(bytes); + EXPECT_EQ(detail::digestPoolLiveCount(), baseline + 1); + + TokenIdExt moved(std::move(source)); // steals the slot — no new slot + EXPECT_EQ(detail::digestPoolLiveCount(), baseline + 1); + EXPECT_EQ(moved.digest(), bytes); + + TokenIdExt target(other); // second live slot + EXPECT_EQ(detail::digestPoolLiveCount(), baseline + 2); + target = std::move(moved); // releases target's slot, adopts moved's + EXPECT_EQ(detail::digestPoolLiveCount(), baseline + 1); + EXPECT_EQ(target.digest(), bytes); + } // source and moved are moved-from and free nothing + EXPECT_EQ(detail::digestPoolLiveCount(), baseline); +}🤖 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 `@cpp/tests/unit_tests/batch_manager/kvCacheManagerV2DigestPoolTest.cpp` around lines 94 - 140, Add tests in DigestPoolTest covering TokenIdExt move construction and move assignment, including the moved-from sentinel behavior. Verify moving transfers digest ownership without changing detail::digestPoolLiveCount(), the moved-from object remains safely destructible without freeing a slot, and move-assigning onto a live digest handle releases the target slot exactly once while preserving the source digest in the destination.cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp (1)
553-565: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winPrevent exceptions from escaping the
EventManagerTestBlockdestructor.
~EventManagerTestBlockcallsclose(), andclose()callsblock->unlinkPage(...). A destructor is implicitlynoexcept, so any exception fromunlinkPageterminates the process instead of surfacing a test failure. TheKvCachedestructor inkvCache.cppalready applies this pattern.♻️ Proposed guard
~EventManagerTestBlock() { - close(); + try + { + close(); + } + catch (...) + { + // Destructors must not propagate exceptions. + } }🤖 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 `@cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp` around lines 553 - 565, Update EventManagerTestBlock::~EventManagerTestBlock to prevent exceptions from escaping during destruction, following the existing KvCache destructor pattern in kvCache.cpp. Guard the close() call so exceptions from block->unlinkPage are handled without triggering destructor-induced process termination, while preserving normal cleanup behavior.cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp (1)
1909-1909: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueConsider passing
textOnly()tomatchReusehere.
matchReusedefaultsknownNoDigesttofalse, so this call always takes the scanning path. The committed tokens of a text-only sequence are known digest-free, somatchReuse(mReuseScope, toSpan(mCommittedTokens), textOnly())would keep the fast path consistent with the commit and partial-snapshot call sites. Behavior is unchanged either way; only hashing cost differs.🤖 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 `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp` at line 1909, Update the matchReuse call in the relevant committed-token reuse path to pass textOnly() as its knownNoDigest argument, matching the commit and partial-snapshot call sites while preserving existing behavior and enabling the digest-free fast path.tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py (1)
309-310: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd the missing
-> Noneannotation.
_KVCache.__init__now has a changed signature, but the implementation still omits its return annotation. The stub intensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyialready declares-> None.As per coding guidelines, annotate every function and use
Nonefor procedures.🤖 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 `@tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py` around lines 309 - 310, Update the _KVCache.__init__ method signature to include the missing -> None return annotation, matching the existing stub declaration and the project’s procedure annotation convention.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.
Inline comments:
In `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.md`:
- Around line 263-267: Update the testing guidance in the AGENTS.md section to
present direct Python file execution only as an optional fast path, remove the
instruction forbidding pytest, and require running the repository validation
command “pytest tests/unittest/” for relevant unit-test changes.
- Around line 218-222: Update the SHA-256 directory reference in the guide’s
hash integration guidance from cpp/tensorrt_llm/batch_manager/common/sha256 to
cpp/tensorrt_llm/common/sha256, keeping the surrounding build and dependency
guidance unchanged.
In `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp`:
- Around line 141-177: Make the validation preceding knownNoDigest=true
release-active so negative int32 token values cannot reach Hasher::update
through the zero-copy withTokens path. Validate every token before marking the
buffer as digest-free, or route any negative-token buffer through the existing
safe conversion path; preserve the normal zero-copy path only when all tokens
are valid non-digest values.
In `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h`:
- Around line 582-583: Update the comment immediately above mTextOnly to
describe it as the resolved per-sequence text-only setting, removing the
inaccurate nullopt/inheritance behavior. Keep the bool declaration and
constructor logic unchanged.
In `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/tokenIdExt.cpp`:
- Around line 259-272: Update TokenIdExt::operator== and the corresponding
operator!= declaration/definition to avoid calling digest() when either digest
handle has the kValueMask sentinel index; preserve the existing equality
behavior for valid digests and let invalid default-constructed or moved-from
handles fail through the intended diagnosable path rather than terminating from
noexcept.
In `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/tokenIdExt.h`:
- Around line 84-90: Update the contiguous int32 token-processing path to
validate every token ID is non-negative before constructing the TokenSpan, using
an always-active check rather than TLLM_CHECK_DEBUG. Ensure invalid negative IDs
are rejected before hashing or any digest-pool handling, while preserving the
existing valid-ID flow and TokenIdExt bounds contract.
In `@cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp`:
- Around line 2130-2138: Update the gen_multimodal_cache_key_tokens binding to
validate that multiModalDataDigest has exactly kv::kDIGEST_LEN bytes before
constructing the vector or calling kv::genMultimodalCacheKeyTokens; reject both
shorter and longer inputs at the binding boundary, preserving the existing
valid-input behavior.
- Around line 113-127: Replace the debug-only validation in the zero-copy branch
of withTokens with unconditional validation that rejects any negative int32
token ID before the reinterpret_cast to kv::TokenIdExt. Preserve the existing
valid-input zero-copy path and error context, while ensuring release builds
cannot treat negative values as digest handles.
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py`:
- Around line 2534-2545: Import _cpp_introspection from
tensorrt_llm.runtime.kv_cache_manager_v2 and retain the existing backend
selection in _reuse_token_source, so the selector is defined when choosing
between get_tokens_view and get_tokens.
In `@tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py`:
- Around line 38-50: Update gen_multimodal_cache_key_tokens to validate that
multi_modal_data_digest has the required backend digest length before generating
tokens, matching the kDIGEST_LEN contract enforced by blockRadixTree.cpp.
Preserve the existing token generation and offset assertions after this
validation.
- Around line 80-98: Update the bulk hashing branch in the block-hashing method
to validate every integer token is in the existing 31-bit range before calling
array("I", data), while preserving the bytes-item fallback behavior. Add a
one-time assertion that the native unsigned-int item size is 4 bytes before
using the packed representation, reusing the existing NDEBUG-aware validation
convention.
In `@tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py`:
- Around line 340-344: Make destructor cleanup safe after partial construction:
in _KVCache.__init__, initialize fields used by __del__ before the text_only
validation can raise; in KVCacheManager.__init__, initialize ownership fields
before StorageManager or BlockRadixTree construction, and guard cleanup in
__del__ so unset fields are tolerated while preserving radix-tree-before-storage
teardown. Apply changes in
tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py (lines 340-344) and
tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py (lines
261-273).
- Around line 340-344: Update the constructor validation and the commit() method
in the KV cache manager to reject digest-token inputs whenever the manager is
configured with text_only=True, before reuse setup or token appending can store
them in _committed_tokens. Add regression coverage for both ingress paths, and
run the salting/reuse tests against both backends.
In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py`:
- Around line 1145-1157: Rename the later duplicate test method
test_v2_kv_cache_event_manager_omits_partial_life_cycle_coverage to a unique,
descriptive name while preserving its assertions and behavior. Ensure both test
definitions are independently collected by pytest, then run the event-manager
tests.
---
Outside diff comments:
In `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp`:
- Around line 586-599: Update removeSubtree in
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp: release
each block with releasePages() before detaching it, revise the lifecycle
comments, and correct “parent.s next map” to “parent’s next map”. Apply the
equivalent change in remove_subtree at
cpp/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py: call
_release_pages() before each detachment and update its stale lifecycle comments.
---
Nitpick comments:
In `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp`:
- Around line 285-297: Remove the redundant static qualifier from the isPrefix
function declaration, leaving its implementation and anonymous-namespace scope
unchanged.
In `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp`:
- Line 1909: Update the matchReuse call in the relevant committed-token reuse
path to pass textOnly() as its knownNoDigest argument, matching the commit and
partial-snapshot call sites while preserving existing behavior and enabling the
digest-free fast path.
In `@cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp`:
- Around line 553-565: Update EventManagerTestBlock::~EventManagerTestBlock to
prevent exceptions from escaping during destruction, following the existing
KvCache destructor pattern in kvCache.cpp. Guard the close() call so exceptions
from block->unlinkPage are handled without triggering destructor-induced process
termination, while preserving normal cleanup behavior.
In `@cpp/tests/unit_tests/batch_manager/kvCacheManagerV2DigestPoolTest.cpp`:
- Around line 94-140: Add tests in DigestPoolTest covering TokenIdExt move
construction and move assignment, including the moved-from sentinel behavior.
Verify moving transfers digest ownership without changing
detail::digestPoolLiveCount(), the moved-from object remains safely destructible
without freeing a slot, and move-assigning onto a live digest handle releases
the target slot exactly once while preserving the source digest in the
destination.
In `@tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py`:
- Line 68: Update the __init__ data parameter annotation so None appears at the
end of the union, preserving the existing int, bytes, and Sequence[int | bytes]
alternatives while resolving Ruff RUF036.
In `@tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py`:
- Around line 309-310: Update the _KVCache.__init__ method signature to include
the missing -> None return annotation, matching the existing stub declaration
and the project’s procedure annotation convention.
In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py`:
- Around line 796-850: Extend test_int32_ndarray_ingest_matches_list with
validation cases for a non-contiguous int32 ndarray view and an unsupported
dtype, asserting both are rejected by the ndarray ingest/probe binding before
buffer reinterpretation. Keep the existing valid contiguous int32 compatibility
checks unchanged, and verify the modified test is registered in the appropriate
test-db/qa list files.
In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py`:
- Around line 54-85: The salting tests do not cover cross-language root-key
parity, large token IDs, or rejection of digest-bearing sequences in the
knownNoDigest/text_only path. Extend the relevant tests around
test_reuse_scope_seeds_distinct_keys and
test_blockchain_keys_are_seeded_by_reuse_scope to compare Python roots with C++
RootBlock.makeKey for identical ReuseScope values, exercise token IDs at or
above 2**31 through both backends, and verify digest-bearing input is rejected
by the knownNoDigest/text_only flow. Run the relevant unit suites with both
TLLM_KV_CACHE_MANAGER_V2_BACKEND=cpp and python.
🪄 Autofix
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: c03a6cb0-57d1-4fe4-a35f-eca94a47d123
📒 Files selected for processing (35)
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.mdcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txtcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/tokenIdExt.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/tokenIdExt.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/math.hcpp/tensorrt_llm/nanobind/batch_manager/bindings.cppcpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cppcpp/tests/unit_tests/batch_manager/CMakeLists.txtcpp/tests/unit_tests/batch_manager/kvCacheManagerV2DigestPoolTest.cppcpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpptensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/runtime/kv_cache_manager_v2/AGENTS.mdtensorrt_llm/runtime/kv_cache_manager_v2/__init__.pytensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyitensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.pytensorrt_llm/runtime/kv_cache_manager_v2/_cache_key.pytensorrt_llm/runtime/kv_cache_manager_v2/_config.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.pytensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.pytensorrt_llm/runtime/kv_cache_manager_v2/_introspection.pytensorrt_llm/runtime/kv_cache_manager_v2/_page.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py
💤 Files with no reviewable changes (3)
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h
- tensorrt_llm/runtime/kv_cache_manager_v2/_cache_key.py
- tensorrt_llm/runtime/kv_cache_manager_v2/_page.py
1 task
lowsfer
force-pushed
the
kvcm2-cpp-follow-up
branch
from
August 5, 2026 12:34
154629d to
cd4681a
Compare
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (4)
cpp/tests/unit_tests/batch_manager/kvCacheManagerV2DigestPoolTest.cpp (1)
211-213: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert reclaimed-slot reuse.
Lines 211-213 only verify that
refillretains its digest value. An allocator that allocates a new tail slot will also pass this test. Record an observable identifier for a released low slot and assert thatrefillreuses it.As per PR objectives, this layer adds lifecycle teardown and digest ownership regression coverage.
🤖 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 `@cpp/tests/unit_tests/batch_manager/kvCacheManagerV2DigestPoolTest.cpp` around lines 211 - 213, Update the test around TokenIdExt refill allocation to retain the observable identifier of a previously released low slot, then assert that the new refill allocation reuses that identifier rather than only checking its digest. Preserve the existing distinctDigest(9999) validation while adding the reclaimed-slot reuse assertion.tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py (1)
400-409: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocument the
text_onlyargument.
create_kv_cache()exposestext_only, but its docstring does not describe it. Add a Google-styleArgs:entry that defines its values, default behavior, and digest-token effect.As per coding guidelines, “Use docstrings rather than comments for externally usable interfaces, Google-style docstrings for classes and functions.”
🤖 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 `@tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py` around lines 400 - 409, Update the create_kv_cache() docstring by adding a Google-style Args entry for text_only. Document the meaning of True, False, and None, including the default behavior and how the setting affects digest-token handling.Source: Coding guidelines
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h (1)
126-182: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse Doxygen for the new public view interfaces.
Document
Span,toSpan, andTokenSpanwith Doxygen comments. State that the source buffer must outlive the view.As per coding guidelines, “document new interfaces with Doxygen.”
🤖 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 `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h` around lines 126 - 182, Add Doxygen comments for the public Span struct, toSpan function, and TokenSpan alias, describing their view semantics and explicitly stating that the referenced source buffer must outlive the view. Keep the existing implementation and ownership behavior unchanged.Source: Coding guidelines
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp (1)
285-295: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueBrace both
ifbodies inisPrefix.Lines 289 and 293 use unbraced
ifbodies. Add braces to comply with the C++ control-flow rule.As per coding guidelines, “always brace if/else, loop, and switch bodies.”
🤖 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 `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp` around lines 285 - 295, Update isPrefix to wrap both if statement bodies— the prefixLen > fullLen early return and the prefix[i] != full[i] return—inside braces, preserving the existing control flow.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.
Inline comments:
In `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h`:
- Around line 135-136: Update the member declarations in the Span<T> definition
to default-initialize ptr to nullptr and len to 0, preserving aggregate
initialization behavior. Ensure default-constructed Span instances are safe to
use before explicit assignment.
---
Nitpick comments:
In `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp`:
- Around line 285-295: Update isPrefix to wrap both if statement bodies— the
prefixLen > fullLen early return and the prefix[i] != full[i] return—inside
braces, preserving the existing control flow.
In `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h`:
- Around line 126-182: Add Doxygen comments for the public Span struct, toSpan
function, and TokenSpan alias, describing their view semantics and explicitly
stating that the referenced source buffer must outlive the view. Keep the
existing implementation and ownership behavior unchanged.
In `@cpp/tests/unit_tests/batch_manager/kvCacheManagerV2DigestPoolTest.cpp`:
- Around line 211-213: Update the test around TokenIdExt refill allocation to
retain the observable identifier of a previously released low slot, then assert
that the new refill allocation reuses that identifier rather than only checking
its digest. Preserve the existing distinctDigest(9999) validation while adding
the reclaimed-slot reuse assertion.
In `@tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py`:
- Around line 400-409: Update the create_kv_cache() docstring by adding a
Google-style Args entry for text_only. Document the meaning of True, False, and
None, including the default behavior and how the setting affects digest-token
handling.
🪄 Autofix
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: 1bbf5146-a312-4457-ba60-9fef63c25bef
📒 Files selected for processing (35)
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.mdcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txtcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/tokenIdExt.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/tokenIdExt.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/math.hcpp/tensorrt_llm/nanobind/batch_manager/bindings.cppcpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cppcpp/tests/unit_tests/batch_manager/CMakeLists.txtcpp/tests/unit_tests/batch_manager/kvCacheManagerV2DigestPoolTest.cppcpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpptensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/runtime/kv_cache_manager_v2/AGENTS.mdtensorrt_llm/runtime/kv_cache_manager_v2/__init__.pytensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyitensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.pytensorrt_llm/runtime/kv_cache_manager_v2/_cache_key.pytensorrt_llm/runtime/kv_cache_manager_v2/_config.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.pytensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.pytensorrt_llm/runtime/kv_cache_manager_v2/_introspection.pytensorrt_llm/runtime/kv_cache_manager_v2/_page.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py
💤 Files with no reviewable changes (3)
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h
- tensorrt_llm/runtime/kv_cache_manager_v2/_page.py
- tensorrt_llm/runtime/kv_cache_manager_v2/_cache_key.py
🚧 Files skipped from review as they are similar to previous changes (28)
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt
- cpp/tests/unit_tests/batch_manager/CMakeLists.txt
- tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.md
- tensorrt_llm/runtime/kv_cache_manager_v2/_config.py
- tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py
- tensorrt_llm/runtime/kv_cache_manager_v2/AGENTS.md
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/math.h
- tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
- cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/tokenIdExt.h
- tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py
- tensorrt_llm/runtime/kv_cache_manager_v2/init.py
- cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp
- tensorrt_llm/runtime/kv_cache_manager_v2/init.pyi
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/tokenIdExt.cpp
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h
- tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
- tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py
- tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py
- tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp
- cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp
Member
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #64054 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #64054 [ run ] completed with state
|
lowsfer
force-pushed
the
kvcm2-cpp-follow-up
branch
from
August 5, 2026 15:07
830292b to
8b533aa
Compare
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Member
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #64069 [ run ] triggered by Bot. Commit: |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h (1)
213-220: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winUse Doxygen comments for the changed public APIs.
commit,textOnly, andsetTextOnlyare public interfaces. Their current comments do not create Doxygen API documentation. Use//!documentation and specify theTokenSpanlifetime contract and text-only validation behavior.As per coding guidelines, “Use
//!and//!<for Doxygen comments, and document new interfaces with Doxygen.”Also applies to: 394-402
🤖 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 `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h` around lines 213 - 220, Convert the public API comments for commit, textOnly, and setTextOnly to Doxygen comments using //! (and //!< where appropriate). Preserve the existing descriptions while explicitly documenting the TokenSpan lifetime requirements for commit and the validation behavior of the text-only interfaces.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.
Inline comments:
In `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h`:
- Around line 173-174: Preserve the digest-free invariant whenever mTextOnly is
true: reject digest-bearing reuse matches and commit spans in the relevant reuse
and commit logic. After _setupForReuse(), validate mCommittedTokens when
construction resolves mTextOnly to true, and ensure setTextOnly(true) maintains
the same invariant. Add coverage for both reuse-match and commit-span rejection
paths.
In `@tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py`:
- Around line 43-45: Add Google-style docstrings to the externally usable Hasher
class, Hasher.update, gen_multimodal_cache_key_tokens, and the related functions
in the diff. Document integer encoding, digest-to-token behavior, reuse-scope
layout, yielded blockchain-key pairs, parameters, return values, and validation
errors; replace the existing line comment in Hasher.update with its docstring.
- Around line 50-53: Replace the assertions for num_tokens and token_offset in
the synthetic-token validation path with explicit ValueError checks. Reject
num_tokens values that are not positive and token_offset values below zero,
while preserving the existing multi_modal_data_digest length validation.
- Around line 465-467: Update shutdown() to enforce that all live _KVCache
instances have closed before invoking storage.destroy(): after
radix_tree.clear() and before storage teardown, assert that _living_kv_caches is
empty, or explicitly close every remaining cache. Preserve orphan-block page
release and prevent StorageManager destruction while any cache still references
an orphan block.
In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py`:
- Line 824: Update the commit_prompt helper to add an explicit type annotation
for tokens that models both accepted inputs: Sequence[TokenIdExt] and a
contiguous int32 ndarray. Keep the function behavior unchanged and place the
annotation on commit_prompt so it satisfies the “annotate every function”
requirement while reflecting the existing call patterns.
---
Outside diff comments:
In `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h`:
- Around line 213-220: Convert the public API comments for commit, textOnly, and
setTextOnly to Doxygen comments using //! (and //!< where appropriate). Preserve
the existing descriptions while explicitly documenting the TokenSpan lifetime
requirements for commit and the validation behavior of the text-only interfaces.
🪄 Autofix
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: 1333cb32-2d57-4e2e-8073-94df8181cf46
📒 Files selected for processing (35)
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.mdcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txtcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/tokenIdExt.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/tokenIdExt.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/math.hcpp/tensorrt_llm/nanobind/batch_manager/bindings.cppcpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cppcpp/tests/unit_tests/batch_manager/CMakeLists.txtcpp/tests/unit_tests/batch_manager/kvCacheManagerV2DigestPoolTest.cppcpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpptensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/runtime/kv_cache_manager_v2/AGENTS.mdtensorrt_llm/runtime/kv_cache_manager_v2/__init__.pytensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyitensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.pytensorrt_llm/runtime/kv_cache_manager_v2/_cache_key.pytensorrt_llm/runtime/kv_cache_manager_v2/_config.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.pytensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.pytensorrt_llm/runtime/kv_cache_manager_v2/_introspection.pytensorrt_llm/runtime/kv_cache_manager_v2/_page.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py
💤 Files with no reviewable changes (3)
- tensorrt_llm/runtime/kv_cache_manager_v2/_page.py
- tensorrt_llm/runtime/kv_cache_manager_v2/_cache_key.py
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h
🚧 Files skipped from review as they are similar to previous changes (27)
- tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py
- tensorrt_llm/runtime/kv_cache_manager_v2/_config.py
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h
- tensorrt_llm/runtime/kv_cache_manager_v2/init.py
- tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/tokenIdExt.cpp
- tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py
- tensorrt_llm/runtime/kv_cache_manager_v2/init.pyi
- tensorrt_llm/runtime/kv_cache_manager_v2/AGENTS.md
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.md
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/math.h
- tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h
- tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py
- tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h
- cpp/tests/unit_tests/batch_manager/CMakeLists.txt
- cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp
- cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/tokenIdExt.h
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp
- cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp
Comment on lines
+173
to
+174
| std::optional<RequestIdType> id, PriorityCb priorityCb, std::optional<int> expectedPromptLength = std::nullopt, | ||
| std::optional<bool> textOnly = std::nullopt); |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Preserve the digest-free invariant for all token ingestion.
The constructor can enable mTextOnly before _setupForReuse. commit then accepts later digest tokens without validation. setTextOnly(true) only scans tokens that already exist in mCommittedTokens.
A cache can therefore contain digest tokens while textOnly() permits key generation to skip digest handling. Reject digest-bearing reuse matches and commit spans when mTextOnly is true. Add coverage for both paths.
Proposed implementation direction
void KvCache::commit(TokenSpan tokens, bool isEnd)
{
+ if (mTextOnly
+ && std::any_of(tokens.begin(), tokens.end(),
+ [](TokenIdExt const& token) { return token.isDigest(); }))
+ {
+ throw std::invalid_argument("Cannot commit digest tokens to a text-only KV cache");
+ }
+
// Existing commit logic.
}Also validate mCommittedTokens after _setupForReuse() when construction resolves mTextOnly to true.
Also applies to: 220-220, 399-402
🤖 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 `@cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h` around lines
173 - 174, Preserve the digest-free invariant whenever mTextOnly is true: reject
digest-bearing reuse matches and commit spans in the relevant reuse and commit
logic. After _setupForReuse(), validate mCommittedTokens when construction
resolves mTextOnly to true, and ensure setTextOnly(true) maintains the same
invariant. Add coverage for both reuse-match and commit-span rejection paths.
Comment on lines
+43
to
+45
| def gen_multimodal_cache_key_tokens( | ||
| id_offset: int, multi_modal_data_digest: bytes, num_tokens: int, token_offset: int = 0 | ||
| ) -> list[TokenIdExt]: |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Add Google-style documentation for the new cache-key API.
Hasher has no class docstring. Hasher.update uses a line comment instead of a docstring. The new functions do not document arguments, return values, and validation errors in Google style.
Document the integer encoding, digest-token behavior, reuse-scope layout, and yielded blockchain-key pairs.
As per coding guidelines, “Use docstrings rather than comments for externally usable interfaces, Google-style docstrings for classes and functions.”
Also applies to: 60-81, 112-132
🤖 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 `@tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py` around lines
43 - 45, Add Google-style docstrings to the externally usable Hasher class,
Hasher.update, gen_multimodal_cache_key_tokens, and the related functions in the
diff. Document integer encoding, digest-to-token behavior, reuse-scope layout,
yielded blockchain-key pairs, parameters, return values, and validation errors;
replace the existing line comment in Hasher.update with its docstring.
Source: Coding guidelines
Comment on lines
+50
to
+53
| if len(multi_modal_data_digest) != _SHA256_DIGEST_SIZE: | ||
| raise ValueError(f"multi_modal_data_digest must have length {_SHA256_DIGEST_SIZE}") | ||
| assert num_tokens > 0 | ||
| assert token_offset >= 0 |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Raise ValueError for invalid synthetic-token arguments.
Assertions can be disabled with Python optimization. Invalid num_tokens or token_offset can then create an empty or malformed synthetic-token sequence. Use explicit ValueError checks, consistent with the existing digest-length validation.
Proposed fix
- assert num_tokens > 0
- assert token_offset >= 0
+ if num_tokens <= 0:
+ raise ValueError("num_tokens must be positive")
+ if token_offset < 0:
+ raise ValueError("token_offset must be non-negative")As per coding guidelines, “raise ValueError rather than assertions.”
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if len(multi_modal_data_digest) != _SHA256_DIGEST_SIZE: | |
| raise ValueError(f"multi_modal_data_digest must have length {_SHA256_DIGEST_SIZE}") | |
| assert num_tokens > 0 | |
| assert token_offset >= 0 | |
| if num_tokens <= 0: | |
| raise ValueError("num_tokens must be positive") | |
| if token_offset < 0: | |
| raise ValueError("token_offset must be non-negative") |
🤖 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 `@tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py` around lines
50 - 53, Replace the assertions for num_tokens and token_offset in the
synthetic-token validation path with explicit ValueError checks. Reject
num_tokens values that are not positive and token_offset values below zero,
while preserving the existing multi_modal_data_digest length validation.
Source: Coding guidelines
Comment on lines
+465
to
+467
| Cleanup is normally deferred to ``__del__``. An orphan block may remain | ||
| referenced by a live ``_KVCache`` and retain its pages until that cache closes; | ||
| every cache must close before ``StorageManager`` teardown. |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py \
--items all --type function
rg -n -C 6 --type py \
'def shutdown\(|def close\(|_living_kv_caches|living_kv_caches|storage\.destroy\(\)' \
tensorrt_llm/runtime/kv_cache_manager_v2
rg -n -C 4 --type py \
'\.shutdown\(' \
tensorrt_llm testsRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
# Find the relevant files and examine shutdown implementation
fd -t f '_kv_cache_manager.py' tensorrt_llm/runtime/kv_cache_manager_v2/
# Look at the specific lines mentioned in the review comment
sed -n '460,470p' tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py
sed -n '620,630p' tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py
# Search for the StorageManager.shutdown() method
rg -n 'def shutdown' tensorrt_llm/runtime/kv_cache_manager_v2/ -A 15
# Look for storage.destroy() calls and their context
rg -n 'storage\.destroy' tensorrt_llm/runtime/kv_cache_manager_v2/ -B 5 -A 2
# Search for _living_kv_caches references
rg -n '_living_kv_caches|living_kv_caches' tensorrt_llm/runtime/kv_cache_manager_v2/ -B 3 -A 3Repository: NVIDIA/TensorRT-LLM
Length of output: 12163
🏁 Script executed:
# Get more context on the shutdown sequence
sed -n '295,315p' tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py
# Look at the Block.__del__ implementation
rg -n '__del__' tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py -A 10
# Check the remove_subtree implementation
rg -n 'def remove_subtree' tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py -A 20
# Get the specific lines mentioned in the review
sed -n '460,475p' tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py
# Check what "orphan block" means in the context
rg -n 'orphan' tensorrt_llm/runtime/kv_cache_manager_v2/ -B 2 -A 2Repository: NVIDIA/TensorRT-LLM
Length of output: 9830
🏁 Script executed:
# Check if there's any guard that prevents shutdown when live caches exist
rg -n 'def shutdown' tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py -A 20 | head -40
# Look at how shutdown is called in LLM API to see the expected sequence
rg -n 'def shutdown' tensorrt_llm/llm.py -A 15
# Check _KVCache.close() to see what it does
rg -n 'def close' tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py -A 10
# Check if there's any validation before storage.destroy()
sed -n '299,320p' tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 2809
🏁 Script executed:
# Find where LLM.shutdown() is defined and how it relates to _kv_cache_manager
fd -t f 'llm.py' tensorrt_llm --max-depth 3
# Look for the actual LLM class
find tensorrt_llm -name '*.py' -type f | xargs grep -l 'class LLM' | head -3
# Check how kv_cache_manager.shutdown() is called
rg -n 'kv_cache_manager\.shutdown|_kv_cache_manager\.shutdown' tensorrt_llm -B 3 -A 3
# See if there's any mechanism to ensure caches are closed before shutdown
rg -n '_living_kv_caches' tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py -B 2 -A 2
# Check if shutdown validates that _living_kv_caches is empty
rg -n 'assert.*_living_kv_caches|if.*_living_kv_caches' tensorrt_llm/runtime/kv_cache_manager_v2/Repository: NVIDIA/TensorRT-LLM
Length of output: 4662
🏁 Script executed:
# Check the LLM API shutdown to see how it manages the cache lifecycle
cat tensorrt_llm/llmapi/llm.py | grep -n 'def shutdown' -A 30 | head -50
# Check if there's a close() method in LLM that closes KV caches first
cat tensorrt_llm/llmapi/llm.py | grep -n 'def close' -A 30 | head -50
# Look for how the test code manages shutdown
rg -n 'llm\.shutdown\(\)' tests/unittest/llmapi/test_llm_pytorch.py -B 5 | head -40
# Check if kv_cache_manager is part of the LLM shutdown sequence
rg -n 'shutdown.*kv_cache|_kv_cache.*shutdown' tensorrt_llm/llmapi/ -B 2 -A 2Repository: NVIDIA/TensorRT-LLM
Length of output: 3101
🏁 Script executed:
# Check what executor.shutdown() does - look for the GenerationExecutor class
rg -n 'class GenerationExecutor' tensorrt_llm --max-depth 5 -A 5
# Find where GenerationExecutor is defined and check its shutdown method
find tensorrt_llm -name '*.py' -type f | xargs grep -l 'class GenerationExecutor' | head -5
# Look for the actual shutdown implementation in executors
rg -n 'def shutdown' tensorrt_llm/_torch/pyexecutor -B 2 -A 20 | head -80
# Check the pyexecutor resources and how kv_cache_manager is cleaned up
rg -n 'kv_cache_manager.*shutdown\|shutdown.*kv_cache' tensorrt_llm/_torch/pyexecutor -B 5 -A 5Repository: NVIDIA/TensorRT-LLM
Length of output: 9759
🏁 Script executed:
# Look at the _util.py where kv_cache_manager.shutdown() is called
sed -n '1920,1950p' tensorrt_llm/_torch/pyexecutor/_util.py
# Check the full cleanup sequence in _util.py
rg -n 'def cleanup_resources\|def shutdown' tensorrt_llm/_torch/pyexecutor/_util.py -A 30 | head -100
# Check if there's any mechanism to close kv caches before the manager shuts down
rg -n 'close\|_living_kv_caches' tensorrt_llm/_torch/pyexecutor/_util.py -B 2 -A 2Repository: NVIDIA/TensorRT-LLM
Length of output: 1518
🏁 Script executed:
# Check how the executor manages cache lifecycle and shutdown
rg -n 'def shutdown' tensorrt_llm/executor/proxy.py -A 30 | head -80
# Look at the worker shutdown to see if it closes caches
rg -n 'def shutdown' tensorrt_llm/executor/worker.py -A 30 | head -80
# Check if there's any cleanup of _living_kv_caches before shutdown
rg -n 'for.*_living_kv_caches\|_living_kv_caches.*close' tensorrt_llm -B 2 -A 2Repository: NVIDIA/TensorRT-LLM
Length of output: 3471
🏁 Script executed:
# Check if there's any validation or guards ensuring caches are closed before shutdown
rg -n 'assert.*_living_kv_caches|if.*_living_kv_caches' tensorrt_llm/runtime/kv_cache_manager_v2/ -B 2 -A 2
# Check how PyExecutor ensures cleanup before calling shutdown
rg -n 'teardown_managers\|shutdown' tensorrt_llm/_torch/pyexecutor/_llm_request_handler.py -B 5 -A 10 | head -60
# Look for any code that explicitly closes caches before calling manager.shutdown()
rg -n 'close.*kv_cache|_living_kv_caches.*close' tensorrt_llm -B 3 -A 3Repository: NVIDIA/TensorRT-LLM
Length of output: 338
🏁 Script executed:
# Find all Python executor files
fd -t f '.*executor.*\.py' tensorrt_llm/_torch/pyexecutor/ | head -20
# Look at the actual resource cleanup and shutdown flow
rg -n 'teardown_managers\|_living_kv_caches' tensorrt_llm/_torch/pyexecutor/ -B 3 -A 3
# Check if Block.__del__ is called and what happens
rg -n 'def __del__' tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py -A 5
# Look for any test that demonstrates the issue
rg -n 'orphan\|_living_kv_caches' tensorrt_llm -B 2 -A 2 | grep -i testRepository: NVIDIA/TensorRT-LLM
Length of output: 992
🏁 Script executed:
# Look at how _KVCache.close() interacts with the tree
sed -n '578,610p' tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
# Check the Block._release_pages() implementation in detail
sed -n '455,480p' tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py
# Look at the documentation on shutdown to see if there's mention of cache cleanup requirements
sed -n '458,475p' tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 3460
Verify that shutdown cannot destroy storage before orphan blocks release pages.
The shutdown() method calls radix_tree.clear() to orphan all blocks, then immediately calls storage.destroy() without checking whether any live _KVCache still holds references to those blocks. The docstring for _release_pages() explicitly requires "every cache must close before StorageManager teardown" (lines 465–467), but shutdown() does not enforce this invariant. An orphan block held by a live _KVCache will not release its pages until the cache closes in __del__, which may occur after storage is destroyed. Add an assertion in shutdown() to verify that _living_kv_caches is empty, or ensure all caches are closed before the storage is destroyed.
🤖 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 `@tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py` around lines
465 - 467, Update shutdown() to enforce that all live _KVCache instances have
closed before invoking storage.destroy(): after radix_tree.clear() and before
storage teardown, assert that _living_kv_caches is empty, or explicitly close
every remaining cache. Preserve orphan-block page release and prevent
StorageManager destruction while any cache still references an orphan block.
| prompt_np = np.asarray(prompt, dtype=np.int32) | ||
| assert prompt_np.dtype == np.int32 and prompt_np.flags["C_CONTIGUOUS"] | ||
|
|
||
| def commit_prompt(tokens) -> None: |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Add a precise type annotation for tokens.
commit_prompt accepts both Sequence[TokenIdExt] and a contiguous int32 ndarray. Its unannotated parameter violates the Python typing rule. Model both accepted input forms in the annotation.
As per coding guidelines, “Annotate every function.”
🤖 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/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py` at line
824, Update the commit_prompt helper to add an explicit type annotation for
tokens that models both accepted inputs: Sequence[TokenIdExt] and a contiguous
int32 ndarray. Keep the function behavior unchanged and place the annotation on
commit_prompt so it satisfies the “annotate every function” requirement while
reflecting the existing call patterns.
Source: Coding guidelines
Replace the temporary Python-to-C++ migration plan with durable guidance colocated with the C++ implementation. Document the architecture, ownership and lifetime model, cache and page state machines, hashing and partial-coverage invariants, storage and CUDA constraints, nanobind boundaries, high-risk changes, and the fast direct-Python test workflow. This preserves useful migration knowledge after the Python backend and its reference documents are removed. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
Replace the variant-based TokenIdExt with a four-byte tagged value. Store normal token IDs inline and keep rare multimodal digests in an address-stable process-global pool with value semantics. Pass digest-free knowledge explicitly to hashing and add manager- and request-level text_only configuration so eligible workloads use bulk SHA-256. Validate the contract at the nanobind boundary to prevent digest tokens from silently corrupting block keys. Move blockchain-key and multimodal-token builders into the native C++ backend while retaining equivalent Python behavior and byte-compatible keys. Remove the obsolete Python cache-key module and cover digest ownership, salting, routing, and statistics with the updated tests. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
…t path Avoid the per-token Python round-trip when feeding request tokens to the KVCacheManagerV2 C++ backend for prefix matching and commit. Tokens were copied to a Python list (get_tokens) then converted back element-by-element (nb::cast per token) into std::vector<TokenIdExt>. A digest-free token is bit-identical across int32_t, the 4-byte TokenIdExt, and the hashed byte stream, so a contiguous int32 buffer is reinterpret_cast to TokenIdExt const* with no copy and no per-token boxing. - Core: TokenSpan (= Span<TokenIdExt const>, moved into common.h) + toSpan() helper; match/createKvCache/probeReuse/matchReuse/KvCache ctor/_setupForReuse/ commit take a TokenSpan (dead match(vector) overload removed). - nanobind: LlmRequest.get_tokens_view(beam) returns a zero-copy read-only int32 ndarray (reference_internal); withTokens() ingest helper takes the int32 fast path (nb::try_cast, convert=false) or the existing per-element fallback for multimodal/digest tokens. - Dispatcher: backend-aware _reuse_token_source() -- the C++ backend gets the zero-copy view, the pure-Python backend keeps a plain list (cannot consume numpy arrays). - Test: int32-ndarray ingest hashes identically to the list path (both directions); skipped on the Python backend. Full KVCacheManagerV2 unit suite green on both backends. Egress get_tokens_view is ~44x faster than get_tokens on a 4K-token prompt (30 us -> 0.7 us), and O(1) in sequence length. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
Validate multimodal digest lengths at the backend boundary and cover digest-token move ownership. Preserve the text-only fast path when matching committed blocks for planned drops. Make partial construction and test-helper teardown safe, clarify orphan block page lifetime, and apply the remaining review-driven naming and documentation cleanups. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
Initialize default Span views safely and document their non-owning lifetime contract. Bring isPrefix in line with control-flow style requirements. Document the per-cache text_only override and strengthen digest-pool coverage to verify reclaimed-slot reuse. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
…2 introspection SsmCommittedPage was folded into CommittedPage, whose num_tokens_in_block now carries the token count for both attention and SSM life cycles. Two introspection helpers still referenced the removed class, breaking the build. - nanobind make_test_block: collapse the two-branch construction into a single CommittedPage. This also fixes the argument order, which passed Priority where the constructor expects numTokensInBlock. - nanobind reuse_match_pages: read page->numTokensInBlock directly instead of down-casting; block storage already holds CommittedPage pointers. Attention pages now report their real token count rather than -1, matching the Python backend, which has always returned it. - Python make_test_block: drop the SsmCommittedPage import, which would have raised ImportError on the pure-Python backend. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
The zero-copy token ingest path calls req.get_tokens_view() whenever the C++ backend is active, but the duck-typed request fakes used by the KVCacheManagerV2 tests only implement get_tokens, so every test driving the block-reuse path failed with AttributeError. Mirror the binding on both fakes: return a contiguous 1-D int32 ndarray, whose dtype is what selects the C++ int32 ingest fast path. _ContextRequest was not reported by CI, since no test list schedules its file, but it drives the same reuse path and would have broken the same way. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
lowsfer
force-pushed
the
kvcm2-cpp-follow-up
branch
from
August 5, 2026 18:13
8b533aa to
0589331
Compare
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Member
Author
|
/bot run --disable-fail-fast |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tests/unittest/_torch/executor/test_kv_cache_manager_v2.py (1)
292-292: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winUse explicit beam-ID validation in both helpers.
assertis disabled bypython -O. A non-defaultbeam_idwould then return the beam-0 token buffer silently.
tests/unittest/_torch/executor/test_kv_cache_manager_v2.py#L292-L292: raiseValueErrorwhenbeam_idis notDEFAULT_BEAM_INDEX.tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py#L94-L94: apply the same explicit validation.As per coding guidelines, raise
ValueErrorrather than using assertions for validation errors.🤖 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_kv_cache_manager_v2.py` at line 292, Replace the assertion validating beam_id with explicit ValueError validation in both helper locations: tests/unittest/_torch/executor/test_kv_cache_manager_v2.py lines 292-292 and tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py lines 94-94. Ensure any beam_id other than DEFAULT_BEAM_INDEX raises ValueError before accessing the beam-0 token buffer.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.
Inline comments:
In `@cpp/tests/unit_tests/batch_manager/kvCacheManagerV2DigestPoolTest.cpp`:
- Around line 218-221: Update the test around the refill object so it verifies
pool cleanup after `refill` is destroyed: retain the current live-count
assertion while `refill` is alive, then leave its scope and assert
`detail::digestPoolLiveCount()` returns to `baseline`. Ensure `tokens` is
cleared before the destruction check as currently intended.
---
Nitpick comments:
In `@tests/unittest/_torch/executor/test_kv_cache_manager_v2.py`:
- Line 292: Replace the assertion validating beam_id with explicit ValueError
validation in both helper locations:
tests/unittest/_torch/executor/test_kv_cache_manager_v2.py lines 292-292 and
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py lines
94-94. Ensure any beam_id other than DEFAULT_BEAM_INDEX raises ValueError before
accessing the beam-0 token buffer.
🪄 Autofix
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: 5b3a9939-d171-4c10-b76c-9ae5f988da3f
📒 Files selected for processing (37)
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.mdcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txtcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/tokenIdExt.cppcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/tokenIdExt.hcpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/math.hcpp/tensorrt_llm/nanobind/batch_manager/bindings.cppcpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cppcpp/tests/unit_tests/batch_manager/CMakeLists.txtcpp/tests/unit_tests/batch_manager/kvCacheManagerV2DigestPoolTest.cppcpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpptensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/runtime/kv_cache_manager_v2/AGENTS.mdtensorrt_llm/runtime/kv_cache_manager_v2/__init__.pytensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyitensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.pytensorrt_llm/runtime/kv_cache_manager_v2/_cache_key.pytensorrt_llm/runtime/kv_cache_manager_v2/_config.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.pytensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.pytensorrt_llm/runtime/kv_cache_manager_v2/_introspection.pytensorrt_llm/runtime/kv_cache_manager_v2/_page.pytests/unittest/_torch/executor/test_kv_cache_manager_v2.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py
💤 Files with no reviewable changes (3)
- tensorrt_llm/runtime/kv_cache_manager_v2/_page.py
- tensorrt_llm/runtime/kv_cache_manager_v2/_cache_key.py
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h
🚧 Files skipped from review as they are similar to previous changes (29)
- tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_salting.py
- cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp
- tensorrt_llm/runtime/kv_cache_manager_v2/AGENTS.md
- tensorrt_llm/runtime/kv_cache_manager_v2/_config.py
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/utils/math.h
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp
- tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/common.h
- tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py
- cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp
- tensorrt_llm/runtime/kv_cache_manager_v2/init.py
- cpp/tests/unit_tests/batch_manager/CMakeLists.txt
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp
- tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py
- tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/CMakeLists.txt
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/tokenIdExt.h
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h
- tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py
- tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py
- tensorrt_llm/runtime/kv_cache_manager_v2/init.pyi
- cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/AGENTS.md
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/tokenIdExt.cpp
- cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp
Comment on lines
+218
to
+221
| tokens.clear(); | ||
| // `refill` still alive here. | ||
| EXPECT_EQ(detail::digestPoolLiveCount(), baseline + 1); | ||
| } |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Assert pool cleanup after refill is destroyed.
Line 220 verifies the live count while refill is still alive. The test exits without verifying that its slot returns to the pool. A leak in the refill allocation path can pass, and later tests establish a new baseline.
Proposed test update
- TokenIdExt const refill(distinctDigest(9999));
- EXPECT_EQ(refill.raw(), reclaimedSlot);
- EXPECT_EQ(refill.digest(), distinctDigest(9999));
-
- tokens.clear();
- // `refill` still alive here.
- EXPECT_EQ(detail::digestPoolLiveCount(), baseline + 1);
+ {
+ TokenIdExt const refill(distinctDigest(9999));
+ EXPECT_EQ(refill.raw(), reclaimedSlot);
+ EXPECT_EQ(refill.digest(), distinctDigest(9999));
+
+ tokens.clear();
+ EXPECT_EQ(detail::digestPoolLiveCount(), baseline + 1);
+ }
+ EXPECT_EQ(detail::digestPoolLiveCount(), baseline);📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| tokens.clear(); | |
| // `refill` still alive here. | |
| EXPECT_EQ(detail::digestPoolLiveCount(), baseline + 1); | |
| } | |
| { | |
| TokenIdExt const refill(distinctDigest(9999)); | |
| EXPECT_EQ(refill.raw(), reclaimedSlot); | |
| EXPECT_EQ(refill.digest(), distinctDigest(9999)); | |
| tokens.clear(); | |
| EXPECT_EQ(detail::digestPoolLiveCount(), baseline + 1); | |
| } | |
| EXPECT_EQ(detail::digestPoolLiveCount(), baseline); |
🤖 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 `@cpp/tests/unit_tests/batch_manager/kvCacheManagerV2DigestPoolTest.cpp` around
lines 218 - 221, Update the test around the refill object so it verifies pool
cleanup after `refill` is destroyed: retain the current live-count assertion
while `refill` is alive, then leave its scope and assert
`detail::digestPoolLiveCount()` returns to `baseline`. Ensure `tokens` is
cleared before the destruction check as currently intended.
Collaborator
|
PR_Github #64091 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #64069 [ 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
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.
Summary
Testing
Dev Engineer Review
KVCacheManagerV2token handling around compact taggedTokenIdExtvalues and pooled multimodal digests.int32token ingestion through nanobind.QA Engineer Review
test_kv_cache_event_manager.pytest_kv_cache_manager_v2.pytest_kv_cache_salting.pykvCacheManagerV2DigestPoolTest.cppkvCacheManagerV2StatsTest.cpp_torch/executor/test_kv_cache_manager_v2.pytest_kv_cache_stats_behavior.pyint32token reuse.tests/integration/test_lists/,test-db/,qa/, orwaives.txtchanges are reported.