[TRTLLM-11160][feat] Clean up SWA work-arounds with the new radix search tree - #12004
[TRTLLM-11160][feat] Clean up SWA work-arounds with the new radix search tree#12004SimengLiu-nv wants to merge 10 commits into
Conversation
|
/bot run --disable-fail-fast |
|
PR_Github #38118 [ run ] triggered by Bot. Commit: |
|
PR_Github #38118 [ run ] completed with state
|
SimengLiu-nv
changed the title
|
/bot run --disable-fail-fast |
|
PR_Github #38157 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #38159 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughThis PR refactors KV cache management in TensorRT-LLM by making queue integrity verification const-correct, adding block lookup tree detection, extending the storeBlocks API with out-of-window block tracking, removing sequence bookkeeping infrastructure (holdSequence/releaseSequence), and replacing subtree detachment with selective per-block detachment from lookup trees. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant WindowBlockManager
participant BlockManager
participant KVCacheBlock
participant LookupTree
Note over Client,LookupTree: Old Flow: Block Release with Subtree Detachment
Client->>WindowBlockManager: release(block)
WindowBlockManager->>WindowBlockManager: freeChildren(block)
loop for each descendant
WindowBlockManager->>KVCacheBlock: mark for eviction
end
WindowBlockManager->>BlockManager: evict marked blocks
Note over Client,LookupTree: New Flow: Block Release with Selective Detachment
Client->>WindowBlockManager: release(block)
alt block in lookup tree
WindowBlockManager->>KVCacheBlock: detachFromLookupNode()
KVCacheBlock->>LookupTree: remove self reference
end
WindowBlockManager->>BlockManager: evict this block only
sequenceDiagram
participant Client
participant WindowBlockManager
participant BlockManager
participant KVCacheBlock
participant EvictionPolicy
Note over Client,EvictionPolicy: storeBlocks with OOW Block Handling
Client->>WindowBlockManager: storeBlocks(blockKeys, blockIds, pinBlocks, numOowBlocks)
loop for each block in storeBlocks
alt is OOW block
alt block has refs OR block in lookup tree
WindowBlockManager->>WindowBlockManager: break (stop storing for reuse)
Note over WindowBlockManager: Conservatively gate reuse on OOW blocks
else OOW block without refs
WindowBlockManager->>KVCacheBlock: set priority = kMinRetentionPriority
WindowBlockManager->>EvictionPolicy: add to eviction queue
end
else regular block
WindowBlockManager->>KVCacheBlock: add to lookup tree
WindowBlockManager->>EvictionPolicy: track for reuse
end
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes 🚥 Pre-merge checks | ✅ 1 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (1 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (1)
cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp (1)
4077-4084: Rename the new helper constants tokUPPER_SNAKE_CASE.
kPENumLayers/kVSWATokensPerBlockdon't match the repository's constant naming convention.As per coding guidelines, "Constants naming should be uppercase snakecase with prefix 'k' (e.g.,
kDIGIT_NUM = 10)."Also applies to: 6394-6401
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp` around lines 4077 - 4084, Rename the newly added constants to the repository's uppercase-snake convention with the leading 'k' prefix and update all usages; e.g., change kPENumLayers -> kPE_NUM_LAYERS, kPENumHeads -> kPE_NUM_HEADS, kPESizePerHead -> kPE_SIZE_PER_HEAD, kPETokensPerBlock -> kPE_TOKENS_PER_BLOCK, kPEMaxNumSequences -> kPE_MAX_NUM_SEQUENCES, kPEBeamWidth -> kPE_BEAM_WIDTH, kPEMaxNewTokens -> kPE_MAX_NEW_TOKENS, and kPEIsStreaming -> kPE_IS_STREAMING (and similarly rename kVSWATokensPerBlock to kVSWA_TOKENS_PER_BLOCK); update every reference to these symbols (also apply the same rename in the other occurrence region mentioned) so compilation and style checks pass.
🤖 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/tensorrt_llm/batch_manager/kvCacheManager.cpp`:
- Around line 1007-1013: The detachFromLookupNode() call mutates the radix-tree
and must be serialized with mCachedBlocksRootMutex to avoid races with
getFreeBlock() and callers like allocateBlock(), replaceSharedBlock(), and
onboardBlock(); wrap block->detachFromLookupNode() (and the related
enqueueRemovedEvent(block, mWindowSize) if it observes tree state) inside a
critical section protected by mCachedBlocksRootMutex or move the logic into a
new helper (e.g., detachBlockFromLookupLocked(block)) that takes
mCachedBlocksRootMutex before calling detachFromLookupNode() so existing callers
need not change.
- Around line 2455-2459: When forcing out-of-window blocks to min priority via
outOfWindowBlock->setPriority(executor::KvCacheRetentionConfig::kMinRetentionPriority),
also clear any retention timeout/expiration so LRUEvictionPolicy::releaseBlock()
won't push the block back into the expiration heap (e.g., reset the block's
durationMs/expiration timer or cancel its retention timeout flag); this prevents
refresh() from later restoring kDefaultPriority via the expiration path and
ensures OOW blocks remain prioritized for eviction.
In `@cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp`:
- Around line 6467-6509: The test VSWAStolenOOWBlockNoCorruption currently only
checks for no crash when seq1 steals an OOW block; update the test to also
verify the stolen block wasn't silently dropped from the trie by, after removing
seq1 (removeSequence(1, llmRequest1)), issuing a new addSequence/addToken
sequence that uses seq1's prefix (e.g., recreate llmRequest1 or a follow-up
request with the same prefix) and assert that block reuse succeeds and
blockManager.getNumFreeBlocks() returns blocksInPrimaryPool afterwards; locate
logic around KVCacheManagerTest::VSWAStolenOOWBlockNoCorruption,
kvCacheManager->addSequence, kvCacheManager->removeSequence and
blockManager.getNumFreeBlocks() to insert the additional follow-up
addSequence/assertion.
- Around line 4103-4351: Tests only check counts/integrity but not which
block/prefix was preserved; update the TruePriorityEviction* tests (e.g.,
TruePriorityEvictionInteriorBlockEvictedFirst,
TruePriorityEvictionHighPriorityInteriorBlockPreserved,
TruePriorityEvictionQueueIntegrityAfterChainEviction,
TruePriorityEvictionNoCrashAfterInteriorEviction) to assert explicit
reuse/non-reuse of the expected prefix blocks instead of just
getNumFreeBlocks()/verifyQueueIntegrity()/getContextCurrentPosition()>0: after
storeContextBlocks/removeSequence, call addSequence/storeContextBlocks for a new
request that shares the critical prefix and assert the exact
contextCurrentPosition equals the expected reused-prefix length (or that a
lookup via kvCacheManager->getBlockManager()/blockManager lookup API returns the
block id for that prefix), and for negative cases assert
contextCurrentPosition==0 or that the original block id is not in the
free/available list; use the existing helpers getBlockManager(), addSequence(),
storeContextBlocks(), removeSequence(), and getContextCurrentPosition() to make
these deterministic assertions so the tests fail if eviction order regresses.
---
Nitpick comments:
In `@cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp`:
- Around line 4077-4084: Rename the newly added constants to the repository's
uppercase-snake convention with the leading 'k' prefix and update all usages;
e.g., change kPENumLayers -> kPE_NUM_LAYERS, kPENumHeads -> kPE_NUM_HEADS,
kPESizePerHead -> kPE_SIZE_PER_HEAD, kPETokensPerBlock -> kPE_TOKENS_PER_BLOCK,
kPEMaxNumSequences -> kPE_MAX_NUM_SEQUENCES, kPEBeamWidth -> kPE_BEAM_WIDTH,
kPEMaxNewTokens -> kPE_MAX_NEW_TOKENS, and kPEIsStreaming -> kPE_IS_STREAMING
(and similarly rename kVSWATokensPerBlock to kVSWA_TOKENS_PER_BLOCK); update
every reference to these symbols (also apply the same rename in the other
occurrence region mentioned) so compilation and style checks pass.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 9ab56718-d8cb-4621-92eb-f36a79736460
📒 Files selected for processing (5)
cpp/include/tensorrt_llm/batch_manager/evictionPolicy.hcpp/include/tensorrt_llm/batch_manager/kvCacheManager.hcpp/tensorrt_llm/batch_manager/evictionPolicy.cppcpp/tensorrt_llm/batch_manager/kvCacheManager.cppcpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp
|
PR_Github #38159 [ run ] completed with state
|
|
/bot run |
|
PR_Github #38322 [ run ] triggered by Bot. Commit: |
Cannot reproduce any of the failures on computelab machines. |
|
PR_Github #38322 [ run ] completed with state
|
|
/bot run |
Infra issue. Unrelated to this PR. |
|
PR_Github #38376 [ run ] triggered by Bot. Commit: |
|
PR_Github #38376 [ run ] completed with state
|
Pod failure, infra issue again. Not related to the code changes. |
|
/bot run |
|
PR_Github #38456 [ run ] triggered by Bot. Commit: |
Signed-off-by: SimengLiu-nv <simengl@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #38681 [ run ] triggered by Bot. Commit: |
|
PR_Github #38681 [ run ] completed with state
|
Signed-off-by: SimengLiu-nv <simengl@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #38767 [ run ] triggered by Bot. Commit: |
Signed-off-by: Simeng Liu <109828133+SimengLiu-nv@users.noreply.github.com>
|
PR_Github #38767 [ run ] completed with state |
Signed-off-by: SimengLiu-nv <simengl@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #38798 [ run ] triggered by Bot. Commit: |
|
Some placeholder related modification covers my work but with difference. My thoughts is make placeholders behave like normal blocks unless necessary so that we don't need to add
|
The PR only uses the placeholder implementation merged in #11919. Feel free to modify placeholder related code in your needs. We can deal with merge conflict depending on which PR lands first. How do you like the plan? |
I agree with dealing conflicts later. Left some comments about trivial changes to make the code cleaner. |
|
PR_Github #38798 [ run ] completed with state
|
One flaky test, 3 node failures. No failure related to the code changes. |
Signed-off-by: SimengLiu-nv <simengl@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #38908 [ run ] triggered by Bot. Commit: |
| // OOW slot: the real block was stored before going OOW (invariant enforced by | ||
| // storeContextBlocks / storeNewBlock). Advance prevBlock via the existing trie | ||
| // value. If the OOW block was evicted (no value at this node), the chain is | ||
| // broken — stop storing subsequent blocks. |
There was a problem hiding this comment.
This is not correct. All blocks that are not placeholders should be stored, unless the tree already has a block for that windowSize and blockKey. There is no such thing as a broken chain. For example, if blocks vector holds [placeholder, placeholder, block1, placeholder, block2], I want block1 and block2 to be stored in the search tree no matter what. I don't care if the blocks that were replaced by placeholders when they went OOW have been evicted or not.
There was a problem hiding this comment.
The reason why we want this behavior is SWA. For example, assume the following blocks are generated for an SWA layer during the lifespan of a sequence (context + generation phases completed):
[b1, b2, b3, b4, b5, b6, b7]
assume the window size is exactly 3 blocks long. The window initially includes [b1,b2,b3], but keeps sliding as new tokens are generated. At some point, b1 goes OOW and is stored in the tree and released to the free queue. When generation completes, the sliding window contains [b4,b5,b6,b7] (if b7 hasn't been filled yet), which means that b1, b2 and b3 have gone OOW at some point and have been released. If there is a lot of demand for KV cache blocks, b1 will be evicted first, since it was the first block to be released. At this point, the search tree contains
[ph, b2, b3]
with your current code, the chain would immediately be broken and no blocks would be stored because b1 has been evicted, which means no reuse can ever happen. Now assume a new sequence is added, which shares a common prefix with the first request that is exactly 5 blocks long. When we look for reusable blocks in loadOrAllocateBlocks, we want to find this if b1 has been evicted:
[ph, b2, b3, b4, b5, ....]
as we see, the last 4 of the first 5 blocks were matched. The first block was not matched because it was evicted, but it does not matter, because the window size is only 3 blocks long, thus we don't need b1 or b2 in order to be able to reuse the first 5 blocks because they are already outside the attention window.
|
PR_Github #38908 [ run ] completed with state
|
|
Does this fix SWA prefix caching reuse? |
|
Close the PR, the changes are migrated to #13346. |
5 participants
Summary by CodeRabbit
Release Notes
Refactor
Tests
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.