[#11879][fix] Fix free-block counter corruption in getFreeBlock offload path#12834
[#11879][fix] Fix free-block counter corruption in getFreeBlock offload path#12834thorjohnsen merged 1 commit intoNVIDIA:mainNVIDIA/TensorRT-LLM:mainfrom eopXD:fix/claim-before-swap-getFreeBlockeopXD/TensorRT-LLM:fix/claim-before-swap-getFreeBlockCopy head branch name to clipboard
Conversation
57936bc to
2e3232b
Compare
📝 WalkthroughWalkthroughThe Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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 1159-1160: Remove the extra space before the inline comment on the
call mEvictionPolicy->claimBlock(offloadBlock); so the comment starts
immediately after the semicolon (i.e., change "claimBlock(offloadBlock); //..."
to "claimBlock(offloadBlock); //...") to satisfy clang-format; update the same
line where mEvictionPolicy->claimBlock(block); and
mEvictionPolicy->claimBlock(offloadBlock); are adjacent to keep consistent
formatting.
🪄 Autofix (Beta)
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: Pro
Run ID: 3f9705ae-f9fb-4ab2-8f53-574314261647
📒 Files selected for processing (1)
cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
|
/bot run --disable-fail-fast |
|
PR_Github #42320 [ run ] triggered by Bot. Commit: |
|
PR_Github #42320 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #42423 [ run ] triggered by Bot. Commit: |
|
PR_Github #42423 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #42624 [ run ] triggered by Bot. Commit: |
|
PR_Github #42624 [ run ] completed with state |
thorjohnsen
left a comment
There was a problem hiding this comment.
This fix is included in #12297, which also fixes the underlying problem in EvictionPolicy class to prevent issue from coming back if somebody in the future decides to move claimBlock calls again. It is an urgent fix, so I am approving it.
|
/bot run --disable-fail-fast |
|
PR_Github #43155 [ run ] triggered by Bot. Commit: |
|
PR_Github #43155 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #43333 [ run ] triggered by Bot. Commit: |
|
PR_Github #43333 [ run ] completed with state
|
94a04dd to
f2a3835
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #43633 [ run ] triggered by Bot. Commit: |
|
PR_Github #43633 [ run ] completed with state
|
f2a3835 to
baf74b7
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #44319 [ run ] triggered by Bot. Commit: |
|
PR_Github #44453 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #44600 [ run ] triggered by Bot. Commit: |
|
PR_Github #44600 [ run ] completed with state
|
baf74b7 to
0af1e55
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #44917 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #45116 [ run ] triggered by Bot. Commit: |
|
PR_Github #45116 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #45323 [ run ] triggered by Bot. Commit: |
|
PR_Github #45323 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #45637 [ run ] triggered by Bot. Commit: |
|
PR_Github #45637 [ run ] completed with state
|
… offload path In WindowBlockManager::getFreeBlock(), the offload-during-allocation path called claimBlock/releaseBlock AFTER swapMemoryPoolBlockOffset, causing getCacheLevel() to return the post-swap level instead of the pre-swap level. This led to erasing from the wrong per-level free queue (undefined behavior) and modifying the wrong level's free-block counter, the root cause of getNumFreeBlocks() exceeding getMaxNumBlocks() and producing negative usedNumBlocks in disaggregated serving prefill mode. Move claims to happen before the swap, matching the correct pattern already used by the standalone offloadBlock() function. Fixes NVIDIA#11879 Signed-off-by: Yueh-Ting Chen <yueh.ting.chen@gmail.com>
0af1e55 to
4e47ba2
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #45886 [ run ] triggered by Bot. Commit: |
|
PR_Github #45886 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #46337 [ run ] triggered by Bot. Commit: |
|
/bot skip --comment "In the latest CI/CD, x86 Single GPU has passed. The failure is not related to our code change. Let us merge the MR without further delay." |
|
PR_Github #46470 [ skip ] triggered by Bot. Commit: |
|
PR_Github #46470 [ skip ] completed with state |
Fundamentally fixes #11879
Summary by CodeRabbit
Description
Root cause of #11879: In WindowBlockManager::getFreeBlock(), when the offload-during-allocation path is taken (primary block evicted to secondary pool to make room), claimBlock()/releaseBlock() were called after
swapMemoryPoolBlockOffset(). This causedgetCacheLevel()— which infers the cache level fromblock->isPrimary()— to return the post-swap level instead of the pre-swap level, resulting in:from the secondary free queue using an iterator into the primary free queue, and vice versa).
causing getNumFreeBlocks(kPrimaryLevel) to exceed getMaxNumBlocks() and usedNumBlocks to go negative.
Fix: Move both claimBlock() calls to happen before the swap, so getCacheLevel() returns the correct pre-swap level. This matches the pattern already used by the standalone WindowBlockManager::offloadBlock() function (line 1307), which correctly claims before swapping.
Before (buggy):
getFreeBlock(primary) → getFreeBlock(secondary) → transfer → swap → claimBlock(wrong level!) → releaseBlock →
claimBlock(wrong level!)
After (fixed):
getFreeBlock(primary) → getFreeBlock(secondary) → claimBlock(primary ✓) → claimBlock(secondary ✓) → transfer → swap →
releaseBlock(secondary ✓)
The final claimBlock(block, priority, durationMs) that follows the offload block becomes a no-op for the queue (the
block's free-iterator is already nullopt) but still correctly applies the caller's retention priority and duration.
Test Coverage
and onboard flows.
path. The linked issue [Bug]: KvEvent Metrics, usedNumBlocks, can have negative block sizes in disagg/prefill mode #11879 was observed in production disagg/prefill deployments. CI run with --disable-fail-fast is
recommended to validate no regressions across disagg test suites.
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.