[#11879][fix] Fix free-block counter corruption in getFreeBlock offload path#12834
[#11879][fix] Fix free-block counter corruption in getFreeBlock offload path#12834eopXD wants to merge 1 commit intoNVIDIA:mainfrom
Conversation
eopXD
force-pushed
the
fix/claim-before-swap-getFreeBlock
branch
from
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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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
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
|
… 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>
eopXD
force-pushed
the
fix/claim-before-swap-getFreeBlock
branch
from
94a04dd to
f2a3835
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #43633 [ run ] triggered by Bot. Commit: |
|
PR_Github #43633 [ run ] completed with state
|
4 participants
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 caused getCacheLevel() — which infers the cache level from block->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.