[None][fix] Emit KV cache removed events for silently detached reuse-tree blocks

[Shang-Pin]

Description

The KV cache event stream (stored/removed events) is the contract external
consumers use to mirror a worker's reuse tree — e.g. a disaggregated KV router
doing prefix-overlap routing. Two paths in WindowBlockManager detach blocks
from the reuse trie without emitting a removed event, so mirrors leak those
hashes forever. This PR fixes both.

Fix 1: blockInRadixTree() tests the wrong thing

WindowBlockManager::blockInRadixTree() uses getPrevBlock() != nullptr as
its "is this block in the reuse trie" test. But getPrevBlock() walks
mLookupNode -> parentNode -> getValue(windowSize) and returns nullptr
whenever the parent node's value slot is empty — which happens as soon as
the parent block is reclaimed. So the check actually answers "does my parent
still hold a block", not "am I in the trie".

Consequence: every descendant of an already-evicted ancestor is misreported as
not-in-tree at its own reclaim. getFreeBlock() and releaseSubtree() guard
their enqueueRemovedEvent() calls with blockInRadixTree(), so those blocks
are silently detached without a removed event.

-    return !block->getUniqueTokens().empty() && block->getPrevBlock() != nullptr;
+    return !block->getUniqueTokens().empty() && block->getLookupNode() != nullptr;

All callers of blockInRadixTree() are removed/updated event-emission
guards and want exactly membership semantics; getLookupNode() is the
canonical membership accessor already used elsewhere in this file.

Fix 2: freeLeafBlock() detaches without an event

WindowBlockManager::freeLeafBlock() (the partial-leaf reuse path in
onboardAndAllocateBlocks) calls block->freeLeafBlock(), which unlinks the
block from the trie directly — bypassing the event-emitting reclaim paths
entirely. Emit a removed event before detaching, under the same
mEventManager && blockInRadixTree(block) guard used by getFreeBlock() and
releaseSubtree().

With Fix 1 alone in production, per-worker instrumentation still showed
detaches split exactly into evicted-with-event + freeLeafBlock-silent
(counter arithmetic: detach_total == evict_emit + free_leaf_silent), i.e.
this was the one remaining silent path.

Impact

Severity scales with how unshared the cached block chains are:

• Disaggregated prefill serving diverse prompts builds deep, unshared
  chains whose ancestors are evicted early. Measured on such a worker: once
  the KV pool began recycling announced blocks, ~97% of reclaims emitted no
  removed event (15,720 suppressed vs 436 emitted in a single 30s window),
  and a downstream KV router's radix tree grew to ~950x the worker's
  physical block capacity (tens of millions of tracked blocks vs ~74K real)
  before OOM.
• Aggregated / decode workloads with heavily shared prefixes keep parent
  chains valued, so removes are emitted normally — which is why this has gone
  unnoticed. Verified: an aggregated gpt-oss-120b worker's event stream had a
  stored:removed block ratio of 0.99.

Verification

• Rolled both fixes across a live disaggregated gpt-oss-120b deployment. At
  the same pool-full point that previously suppressed 15,720 removes/30s, the
  fixed build emitted 11,385 with zero suppression; the publisher's cumulative
  stored-minus-removed block count settled at ~the physical pool size (mirror
  invariant restored), and the downstream KV router's tracked-block count went
  from ~9.4M/hour unbounded growth to flat.
• A/B control: a worker image identical except for reverting these two
  fixes (engine .so swap only) was placed back in the same fleet — the
  router's per-worker tracked-block count immediately resumed unbounded
  growth, while fixed workers stayed flat. The regression is attributable to
  exactly these two code paths.

Test Coverage

The change is two guarded enqueueRemovedEvent() emissions on existing
reclaim paths; existing KV cache manager event-stream unit tests cover the
emission plumbing. Happy to add targeted unit tests asserting (a) a descendant
of an evicted ancestor emits removed on reclaim and (b) freeLeafBlock
emits removed, if maintainers can point at the preferred fixture for
constructing a partial-leaf reuse state.

PR Checklist

• Please check this after reviewing the above items as appropriate for this PR.

Summary by CodeRabbit

• Bug Fixes
  • Improved KV cache block reclamation to correctly emit removal events when cached blocks are released.
  • Prevented stale cache references during partial-leaf reuse scenarios.
  • Improved detection of cached blocks whose parent blocks have already been reclaimed.

[coderabbitai]

📝 Walkthrough

Walkthrough

WindowBlockManager now emits KV cache removal events before freeing radix-tree leaf blocks and determines radix-tree membership through lookup nodes rather than previous-block links.

Changes

KV cache event handling

Layer / File(s)	Summary
Leaf reclamation and radix membership cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp	freeLeafBlock queues a guarded removal event before freeing blocks in the radix tree, and blockInRadixTree checks getLookupNode() alongside non-empty tokens.

Estimated code review effort: 1 (Trivial) | ~5 minutes

Suggested reviewers: chienchunhung, thorjohnsen, athena-nv, brb-nv, eopxd

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title matches the main change and follows the repository's required [None][type] summary format.
Description check	✅ Passed	The description clearly explains the issue, fix, impact, verification, and test coverage, with only minor template omissions.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 1

🤖 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/kvCacheManager.cpp`:
- Around line 1233-1243: Remove the newly added enqueueRemovedEvent call from
releaseSubtree’s block-detachment path to prevent duplicate removal events, and
retain a single emission in freeLeafBlock before the block is detached. Preserve
the existing event behavior for blocks removed through releaseSubtree and other
callers.

🪄 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: Enterprise

Run ID: 8a5ffdd3-282f-4f71-bef6-1856cd749204

📥 Commits

Reviewing files that changed from the base of the PR and between 0c2952e and 526eb71.

📒 Files selected for processing (1)

• cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp

[nvxuanyuc]

LGTM. This PR:

1. Fixes silent block detaches without a removed event emit.
2. Corrects the radix tree membership check.
   Both case tests need new plumbing of event manager in current unit testing fixtures, should be low-effort.

[nvxuanyuc]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #59276 [ run ] triggered by Bot. Commit: 1b1e59e Link to invocation

[SimengLiu-nv]

@Shang-Pin Would you like to add test coverage for the code changes?

[SimengLiu-nv]

Code changes looks good to me. Request coverage tests to be added.

[Shang-Pin]

@SimengLiu-nv Added in 359157c — two regression tests in kvCacheManagerTest.cpp:

• KVCacheManagerEventRemovedForBlocksWithEvictedAncestors (fix 1): uses per-token-range retention priorities to make the eviction policy reclaim a chain head before its descendant, then asserts every hash announced via Stored is un-announced via Removed when its block is reclaimed.
• KVCacheManagerEventRemovedOnPartialLeafReuse (fix 2): a request partially matching a stored leaf triggers the freeLeafBlock detach in onboardAndAllocateBlocks; asserts a Removed event is emitted for the leaf's announced hash.

Validated both directions by compiling the test binary once and running it against a pre-fix and a post-fix libtensorrt_llm.so:

engine	EvictedAncestors	PartialLeafReuse
without this PR	FAILED ("block announced via Stored was reclaimed without a Removed event")	FAILED ("partial-leaf reuse must emit a Removed event for the detached leaf")
with this PR	OK	OK

[tensorrt-cicd]

PR_Github #59276 [ run ] completed with state SUCCESS. Commit: 1b1e59e
/LLM/main/L0_MergeRequest_PR pipeline #47763 completed with status: 'FAILURE'

CI Report

⚠️ 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

CI Agent Failure Analysis

Link to invocation

[SimengLiu-nv]

LGTM.

[SimengLiu-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #59477 [ run ] triggered by Bot. Commit: c4e1752 Link to invocation

[tensorrt-cicd]

PR_Github #59477 [ run ] completed with state FAILURE. Commit: c4e1752
/LLM/main/L0_MergeRequest_PR pipeline #47942 completed with status: 'FAILURE'

CI Report

⚠️ 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

CI Agent Failure Analysis

Link to invocation