[TRTLLM-11159][feat] Wire KVCacheBlock to UnifiedBlockTree, replacing mPrevBlock/mNextBlocks with lookup-node pointers.

[SimengLiu-nv]

Summary by CodeRabbit

Release Notes

• New Features

  • Enhanced KV-cache block management with improved lookup and sharing mechanisms
  • Added support for multimodal data integration in cache blocks

• Bug Fixes

  • Improved cache block lifecycle and pruning efficiency

• Tests

  • Added comprehensive unit tests for cache block tree operations and lookup functionality

Description

Test Coverage

PR 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)

• 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.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[SimengLiu-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #37747 [ run ] triggered by Bot. Commit: 6e37259 Link to invocation

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR introduces a radix tree-based KV cache block management system for TensorRT LLM, consisting of a templated trie data structure, block key abstractions with multimodal support, a specialized radix block tree for cache block lookups, and integration into the existing KV cache manager's block linking infrastructure.

Changes

Cohort / File(s)	Summary
Core Trie Infrastructure cpp/include/tensorrt_llm/batch_manager/templatedTrie.h, cpp/include/tensorrt_llm/batch_manager/stringSetTrie.h	Introduces a generic, templated radix tree (Trie and Node classes) with support for partial matching, multi-key values, and node lifecycle management. Adds StringSet, a specialized trie for string prefix operations with insert, erase, and contains methods.
Block Key Management cpp/include/tensorrt_llm/batch_manager/blockKey.h, cpp/tensorrt_llm/batch_manager/blockKey.cpp	Defines BlockKey struct for cache block identification with support for multimodal data, LoRA task IDs, and cache salt. Implements block key comparison, shortening, partial matching, and hashing via BlockKeyHasher. Provides helper functions generateBlockHashExtraKeys and buildBlockKeys for multimodal hash key generation.
Radix Block Tree cpp/include/tensorrt_llm/batch_manager/radixBlockTree.h	Specializes the templated Trie for KV cache blocks, providing UnifiedBlockTree with insertBlock and lookupBlock methods for managing cache blocks indexed by token prefix and window size. Defines BlockPtr, LookupNode, and related type aliases.
Hash Utilities cpp/include/tensorrt_llm/batch_manager/common.h	Adds inline hash mixing helpers hash32Mix and hash64Mix in the kv_cache_manager namespace to support non-cryptographic hashing for cache operations.
KV Cache Manager Integration cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h, cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp	Integrates radix tree-based block linking by adding attachToLookupNode, detachFromLookupNode, and setAsRoot methods to KVCacheBlock. Changes getPrevBlock signature to return BlockPtr by value. Adds UnifiedBlockTree member to WindowBlockManager and updates block lifecycle management to support tree-based navigation. Removes legacy BlockKey infrastructure from cache manager.
Build Configuration cpp/tensorrt_llm/batch_manager/CMakeLists.txt	Adds blockKey.cpp to batch manager source compilation targets.
Unit Tests cpp/tests/unit_tests/batch_manager/blockKeyTest.cpp, cpp/tests/unit_tests/batch_manager/radixTreeTest.cpp, cpp/tests/unit_tests/batch_manager/radixBlockTreeTest.cpp, cpp/tests/unit_tests/batch_manager/CMakeLists.txt	Expands test coverage for BlockKey operations (equality, shortening, hashing, partial matching), templated trie functionality (node insertion/lookup, partial matching, value management), and radix block tree lifecycle (attach/detach, traversal, pruning behavior). Updates CMake to register new test targets and copyright year to 2026.

Sequence Diagram(s)

sequenceDiagram
    participant Req as LLM Request
    participant BM as Block Manager
    participant RBT as Radix<br/>Block Tree
    participant BK as BlockKey
    participant KB as KV Cache Block

    Req->>BM: Process request with tokens
    BM->>BK: generateBlockHashExtraKeys(request)
    BK-->>BM: Vector of MmKey entries
    BM->>BK: buildBlockKeys(blocked tokens)
    BK-->>BM: Vector of BlockKey instances
    BM->>RBT: insertBlock(BlockKey prefix, windowSize)
    RBT->>RBT: Create/reuse trie nodes
    RBT->>KB: Store BlockPtr in tree node
    RBT-->>BM: Insertion complete
    
    Note over BM: Later lookup
    BM->>RBT: lookupBlock(BlockKey prefix, windowSize)
    RBT->>RBT: Traverse trie by token prefix
    RBT->>RBT: Match partial keys if needed
    RBT-->>BM: Optional BlockPtr
    BM->>KB: attachToLookupNode(lookup node)
    KB-->>BM: Block linked to tree

Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~75 minutes

🚥 Pre-merge checks | ✅ 1 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 37.41% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Description check	⚠️ Warning	The PR description is incomplete, containing only the repository template with no actual details about the feature, rationale, or implementation approach.	Add a clear description explaining the purpose of wiring KVCacheBlock to UnifiedBlockTree, why this change improves the system, and what the key architectural changes are.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly describes the main feature: wiring KVCacheBlock to UnifiedBlockTree and replacing mPrevBlock/mNextBlocks with lookup-node pointers, which aligns with the primary architectural changes throughout the changeset.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

Tip

Try Coding Plans. Let us write the prompt for your AI agent so you can ship faster (with fewer bugs).
Share your feedback on Discord.

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

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 (2)

cpp/include/tensorrt_llm/batch_manager/common.h (1)

1-2: ⚠️ Potential issue | 🟡 Minor

Update the header year range for this modified file.

This header now has 2026 changes, but the copyright range still ends at 2024.

As per coding guidelines: "All TensorRT-LLM source files should contain an NVIDIA copyright header with the year of the latest meaningful modification."

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@cpp/include/tensorrt_llm/batch_manager/common.h` around lines 1 - 2, Update
the copyright header at the top of this file (the file-level comment in
cpp/include/tensorrt_llm/batch_manager/common.h) so the year range includes the
latest modification year 2026; replace the existing "2023-2024" range with
"2023-2026" (or otherwise ensure 2026 is the ending year) to comply with the
project header guideline.

cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h (1)

1-2: ⚠️ Potential issue | 🟡 Minor

Update the header year range for this modified file.

This header contains 2026 changes, but the copyright range still ends at 2024.

As per coding guidelines: "All TensorRT-LLM source files should contain an NVIDIA copyright header with the year of the latest meaningful modification."

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h` around lines 1 - 2,
Update the copyright header in kvCacheManager.h: replace the existing "Copyright
(c) 2022-2024, NVIDIA CORPORATION." comment block with a header that includes
the latest modification year (e.g., "2022-2026" or "2026" per project
convention) so the range reflects the 2026 changes; ensure the change is made in
the top-of-file header comment where the current copyright line appears.

🧹 Nitpick comments (1)

cpp/tests/unit_tests/batch_manager/blockKeyTest.cpp (1)

164-170: Add a regression test for multimodal start-offset semantics.

These additions cover extraKeys hash differences, but not the case where multimodal content starts inside a block vs before a block. A focused test on generateBlockHashExtraKeys(...).startOffset would lock this behavior down.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@cpp/tests/unit_tests/batch_manager/blockKeyTest.cpp` around lines 164 - 170,
Add a unit test (e.g., TEST_F(BlockKeyTest, MultimodalStartOffsetSemantics))
that constructs two BlockKey instances using makeMmKey with the same multimodal
token but different start positions (one where the multimodal content starts
before the block and one where it starts inside the block), call
generateBlockHashExtraKeys(...) on each and assert the produced .startOffset
values differ and that BlockKeyHasher::hash(...) reflects that difference (i.e.,
hashes are not equal); use the existing helpers BlockKey, BlockKeyHasher,
makeMmKey and generateBlockHashExtraKeys to locate the code to modify.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@cpp/include/tensorrt_llm/batch_manager/blockKey.h`:
- Around line 96-97: The precheck in numMatchingTokens() omits usesExtraIds so
keys with different usesExtraIds can proceed and yield false matches; update the
equality check that currently compares loraTaskId, extraKeys, and cacheSaltID to
also compare usesExtraIds (i.e., ensure usesExtraIds == other.usesExtraIds)
before running the token-prefix logic in numMatchingTokens(), so the function
returns non-matching for keys that differ in usesExtraIds.

In `@cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h`:
- Around line 208-214: The two docblocks are inconsistent about getPrevBlock()
behavior: update the comments for the dummy root wiring (the block documenting
mCachedBlocksRoot and the later doc around lines 239-243) to state the same
contract: that after initialising the dummy root lookup-node link (the routine
described in the root block comment) direct children of the root will return the
root block from getPrevBlock(); before that initialization getPrevBlock() will
return nullptr. Reference getPrevBlock(), mCachedBlocksRoot, the dummy root
lookup-node initialization description and WindowBlockManager in both places and
make the wording symmetric so both docblocks describe the same initialized vs
uninitialized behavior.

In `@cpp/include/tensorrt_llm/batch_manager/radixBlockTree.h`:
- Around line 84-90: lookupBlock currently returns the first (shallowest) valid
vm from valueMatches.matches because it iterates forward; change it to select
the most specific (deepest) valid match by scanning valueMatches.matches from
the end toward the front (or otherwise choose the last valid entry) before
returning vm.value. Update the logic in lookupBlock to examine
valueMatches.matches in reverse order and return the first encountered entry
where vm.isValid && vm.value, leaving lookupValues, allowPartialMatch, and
windowSize usage unchanged.

In `@cpp/include/tensorrt_llm/batch_manager/templatedTrie.h`:
- Around line 401-411: The current _getEdges implementation stops recursion when
mValue is non-empty, so descendant edges are dropped; change the logic in
_getEdges/getEdges so that after adding the current node's edge when mValue is
not empty (using edges.emplace_back(edge)), the method continues to iterate
mNextNodes and call node->_getEdges(edge, edges) for each child; ensure mValue,
edges, _getEdges and mNextNodes are referenced to locate and update the
function.

In `@cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp`:
- Around line 133-155: The attachToLookupNode and detachFromLookupNode functions
currently ignore the boolean results from node->setValue(...) and
mLookupNode->clearValue(...), which lets local members (mLookupNode,
mWindowSize) become out-of-sync with the trie; update attachToLookupNode to
check the return value of node->setValue(windowSize, std::move(self),
/*overwrite=*/false) and only assign mLookupNode and mWindowSize when setValue
returns true (otherwise log or propagate an error), and update
detachFromLookupNode to only reset mLookupNode and mWindowSize when the prior
mLookupNode->clearValue(mWindowSize) returns true (otherwise leave state intact
and log/handle the failure or retry); reference these exact symbols
(attachToLookupNode, detachFromLookupNode, node->setValue,
mLookupNode->clearValue, mLookupNode, mWindowSize) when making the change.

---

Outside diff comments:
In `@cpp/include/tensorrt_llm/batch_manager/common.h`:
- Around line 1-2: Update the copyright header at the top of this file (the
file-level comment in cpp/include/tensorrt_llm/batch_manager/common.h) so the
year range includes the latest modification year 2026; replace the existing
"2023-2024" range with "2023-2026" (or otherwise ensure 2026 is the ending year)
to comply with the project header guideline.

In `@cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h`:
- Around line 1-2: Update the copyright header in kvCacheManager.h: replace the
existing "Copyright (c) 2022-2024, NVIDIA CORPORATION." comment block with a
header that includes the latest modification year (e.g., "2022-2026" or "2026"
per project convention) so the range reflects the 2026 changes; ensure the
change is made in the top-of-file header comment where the current copyright
line appears.

---

Nitpick comments:
In `@cpp/tests/unit_tests/batch_manager/blockKeyTest.cpp`:
- Around line 164-170: Add a unit test (e.g., TEST_F(BlockKeyTest,
MultimodalStartOffsetSemantics)) that constructs two BlockKey instances using
makeMmKey with the same multimodal token but different start positions (one
where the multimodal content starts before the block and one where it starts
inside the block), call generateBlockHashExtraKeys(...) on each and assert the
produced .startOffset values differ and that BlockKeyHasher::hash(...) reflects
that difference (i.e., hashes are not equal); use the existing helpers BlockKey,
BlockKeyHasher, makeMmKey and generateBlockHashExtraKeys to locate the code to
modify.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: c5eb5f0b-720d-41de-b74f-fda5a3cee44a

📥 Commits

Reviewing files that changed from the base of the PR and between 5b81307 and 6e37259.

📒 Files selected for processing (13)

• cpp/include/tensorrt_llm/batch_manager/blockKey.h
• cpp/include/tensorrt_llm/batch_manager/common.h
• cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h
• cpp/include/tensorrt_llm/batch_manager/radixBlockTree.h
• cpp/include/tensorrt_llm/batch_manager/stringSetTrie.h
• cpp/include/tensorrt_llm/batch_manager/templatedTrie.h
• cpp/tensorrt_llm/batch_manager/CMakeLists.txt
• cpp/tensorrt_llm/batch_manager/blockKey.cpp
• cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
• cpp/tests/unit_tests/batch_manager/CMakeLists.txt
• cpp/tests/unit_tests/batch_manager/blockKeyTest.cpp
• cpp/tests/unit_tests/batch_manager/radixBlockTreeTest.cpp
• cpp/tests/unit_tests/batch_manager/radixTreeTest.cpp

[thorjohnsen]

This PR replaces the old search structure with the new constant radix search tree from templatedTrie.h, and it does it in a way that preserves the behavior of the old WindowBlockManager and KVCacheBlock APIs, thus none of the existing unit- and integration tests should break. It opens up the possibility of loosening the block eviction restrictions gradually. For example, we can keep the rule that evicts all children of a newly evicted block for full attention layers and remove it for sliding window attention layers. Evicting children makes sense for full attention layers, because a block cannot be reused if one of it's ancestors have been evicted, but it does not make sense for sliding window attention layers because the evicted ancestor block likely will be outside current attention window and thus is not needed for reuse. Very good initiative Simeng!

I found one or more bugs that must be fixed before merge (see follow-up review).

[tensorrt-cicd]

PR_Github #37747 [ run ] completed with state SUCCESS. Commit: 6e37259
/LLM/main/L0_MergeRequest_PR pipeline #29217 completed with status: 'FAILURE'

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[eopXD]

1. Having the utilities attachToLookupNode, detachFromLookupNode, setAsRoot with LookupNode as the caller will be more natural.
2. Mutex is removed, how do we guarantee a safe read to a node now?
3. I think its better to do a unified search tree in this merge request rather than defer to a "PR B".
4. Please evaluate and see if we need those [[maybe_unused]] variables. For instance, you can return them in its parent function and log them.
5. Update copyright header for the edited files
6. Its a good time to implement freeDescendantsRecursively with a stack. Calling it recursively may impose stack overflow.

[eopXD]

Trie nodes are CPU heap objects not tracked by the block pool. In the worst case the number of live trie nodes is block_pool_size × max_depth, which maybe pretty large.

We should be aware of this, and possibly add logging of this in follow-up work.

[SimengLiu-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #37923 [ run ] triggered by Bot. Commit: 6be60e0 Link to invocation

[SimengLiu-nv]

@eopXD

1. Having the utilities attachToLookupNode, detachFromLookupNode, setAsRoot with LookupNode as the caller will be more natural.

Prefer to keep them in KVCacheBlock and they are "operations" of the block and use/update private variable of KVCacheBlock.

2. Mutex is removed, how do we guarantee a safe read to a node now?

The mCachedBlocksRootMutex should server the purpose.

3. I think its better to do a unified search tree in this merge request rather than defer to a "PR B".

Implemented in the second commit.

4. Please evaluate and see if we need those [[maybe_unused]] variables. For instance, you can return them in its parent function and log them.

Updated in the second commit.

5. Update copyright header for the edited files

Updated in the second commit.

6. Its a good time to implement freeDescendantsRecursively with a stack. Calling it recursively may impose stack overflow.

Implemented in the second commit.

[thorjohnsen]

Great work, just have a couple of questions.

[tensorrt-cicd]

PR_Github #37923 [ run ] completed with state SUCCESS. Commit: 6be60e0
/LLM/main/L0_MergeRequest_PR pipeline #29367 completed with status: 'FAILURE'

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[SimengLiu-nv]

Having the utilities attachToLookupNode, detachFromLookupNode, setAsRoot with LookupNode as the caller will be more natural.
Prefer to keep them in KVCacheBlock and they are "operations" of the block and use/update private variable of KVCacheBlock.

Please find a way to eliminate the std::shared_ptr<KVCacheBlock> self parameter.

Updated in the new commit.

Mutex is removed, how do we guarantee a safe read to a node now?
The mCachedBlocksRootMutex should server the purpose.

For instance, does access to optBlock get a protected read under findMatchingBlock?

The calls are guarded by mCachedBlocksRootMutex or sequential. Will keep a mind to check for potential issue about the mutex removal.

Please evaluate and see if we need those [[maybe_unused]] variables. For instance, you can return them in its parent function and log them.
Updated in the second commit.

I think you still missed some [[maybe_unused]] variables. May you check again?

Nice call, fixed in the new commit.

[SimengLiu-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #38067 [ run ] triggered by Bot. Commit: e3d11fa Link to invocation

[thorjohnsen]

LGTM

[tensorrt-cicd]

PR_Github #38067 [ run ] completed with state SUCCESS. Commit: e3d11fa
/LLM/main/L0_MergeRequest_PR pipeline #29494 completed with status: 'FAILURE'

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[SimengLiu-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #38108 [ run ] triggered by Bot. Commit: 2b48077 Link to invocation

[tensorrt-cicd]

PR_Github #38108 [ run ] completed with state SUCCESS. Commit: 2b48077
/LLM/main/L0_MergeRequest_PR pipeline #29519 completed with status: 'FAILURE'

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

[SimengLiu-nv]

/bot run

[tensorrt-cicd]

PR_Github #38116 [ run ] triggered by Bot. Commit: 2b48077 Link to invocation

[tensorrt-cicd]

PR_Github #38116 [ run ] completed with state SUCCESS. Commit: 2b48077
/LLM/main/L0_MergeRequest_PR pipeline #29527 completed with status: 'SUCCESS'

Link to invocation