[LiRunGuo]

[#17908][fix] Make BlockRadixTree.clear() terminate on a childless RootBlock

Closes #17908

Background

BlockRadixTree.clear() (in tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py)
shrinks self.next only as a side effect of the inner loop: detach_next() prunes a
RootBlock when its last child is detached. A RootBlock that is already childless has
no child to detach, so the outer while self.next: loop re-selects the same root forever
and clear() never returns. The trailing assert not self.next is unreachable in that state.

A childless root is not exotic: RootBlock.__init__ ends with prev.next[self.key] = self,
so add_or_get_existing() publishes an empty root by construction before its caller
attaches a block; any interruption between those two steps leaves the tree permanently
unclearable. The only production caller is KVCacheManager.shutdown(), which therefore
hangs whenever the tree holds a childless root.

tree = BlockRadixTree(object(), 32)
tree.add_or_get_existing(ReuseScope(lora_id=0, salt=None))
tree.clear()   # never returns (pins one core at 100%, no GPU activity)

Summary

Root-edge bookkeeping is centralised into a single new helper,
detach_root_block(root), that removes a RootBlock from its tree's top-level next
registry:

• detach_next() now calls it when the last child of a RootBlock is detached
  (eager prune path — unchanged behaviour, same bookkeeping).
• clear() now calls it directly for a root that never gained a child, so every
  iteration of the loop makes explicit progress (detach a childless root, or remove one
  subtree) and the loop provably terminates.

Changes:

• tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py
  • add detach_root_block();
  • route detach_next()'s root pruning through it;
  • make BlockRadixTree.clear() detach childless roots explicitly.
• tests/unittest/kv_cache_manager_v2_tests/test_block_radix_tree.py (new)
  • 6 CPU-only regression tests (pytest.mark.cpu_only): empty tree, childless-root
    regression, single chain removal (incl. orphan checks), mixed childless + populated
    roots, multi-branch subtrees, idempotence and post-clear reuse of the same reuse scope.

Impacts

• clear() terminates for every reachable tree state; KVCacheManager.shutdown() can no
  longer hang behind a childless root.
• Removal order for populated trees is unchanged (still post-order, leaves first, via the
  existing remove_subtree path), so page-cleanup semantics of Block.__del__ and
  event-manager notifications are untouched.
• No API change; internal only. The fix is pure Python and was validated with CPU-only
  tests in the kv_cache_manager_v2 fast mode
  (TLLM_KV_CACHE_MANAGER_V2_BACKEND=python).

Verification

Reproduced the hang on main (777a2c7) with the snippet above (process does not return;
killed by timeout). After the fix the snippet returns immediately and the new test file
passes; pre-fix, the same tests hang at the first childless-root case.

# fast mode (see kv_cache_manager_v2/AGENTS.md)
PYTHONPATH="$REPO_ROOT/tensorrt_llm/runtime/" \
  python "$REPO_ROOT/tests/unittest/kv_cache_manager_v2_tests/test_block_radix_tree.py" -v

Existing related tests still pass: test_kv_cache_salting.py (3 tests).

Dev Engineer Review

• BlockRadixTree.clear() now removes childless roots and always makes progress.
• detach_root_block() centralizes root cleanup for explicit clearing and eager pruning.
• The change preserves populated-tree removal order and does not change public APIs.

QA Engineer Review

• Modified tests/unittest/kv_cache_manager_v2_tests/test_block_radix_tree.py.
• Added six CPU-only tests for empty, childless, populated, mixed, branched, idempotent, and reusable tree states.
• The tests verify root removal, orphaned blocks, termination, and reuse after clearing.
• The suite is listed in tests/integration/test_lists/test-db/l0_cpu.yml and other applicable test lists.
• Coverage verdict: sufficient.

Per-File QA Perspective

• tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py: Verify childless-root removal, populated-subtree cleanup, eager pruning, and repeated clearing. No API or default behavior changes are introduced.
• tests/unittest/kv_cache_manager_v2_tests/test_block_radix_tree.py: Covers BlockRadixTree.clear() lifecycle and regression paths with CPU-only tests. The containing suite is listed in the CPU CI test list.

[LiRunGuo]

/bot run

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 852385e7-dd9f-4fac-808d-c719e14c74a2

📥 Commits

Reviewing files that changed from the base of the PR and between 0369d45 and 917cbfc.

📒 Files selected for processing (2)

• tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py
• tests/unittest/kv_cache_manager_v2_tests/test_block_radix_tree.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

Walkthrough

The change centralizes root detachment and updates BlockRadixTree.clear() to remove childless roots. CPU-only tests cover clearing, subtree removal, idempotence, orphaning, and tree reuse.

Changes

Root clear lifecycle

Layer / File(s)	Summary
Centralized root detachment tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py	Adds detach_root_block and uses it when detach_next removes a root’s last child.
Clear lifecycle and regression coverage tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py, tests/unittest/kv_cache_manager_v2_tests/test_block_radix_tree.py	clear() explicitly detaches childless roots. CPU-only tests cover empty, chained, branched, mixed, repeated, orphaned, and reused trees.

Priority: ➖ Normal — Schedule the BlockRadixTree cleanup fix because a childless root can make KVCacheManager shutdown hang, while the change remains an internal, narrowly scoped lifecycle correction.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Severity of issue fixed: Medium

Merge Risk: ⚪ Minimal · up to 917cb

This change prevents shutdown hangs caused by childless cache-tree roots and preserves tree cleanup behavior across populated, mixed, and reused trees. The change is ready to merge.

Suggested reviewers: brnguyen2

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 13.33% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 15 functions across 2 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies issue #17908 and the primary fix: making BlockRadixTree.clear() terminate for childless RootBlock instances.
Description check	✅ Passed	The description explains the defect, root cause, implementation, impact, test coverage, and verification. It does not reproduce the template's explicit PR Checklist section, but the required technical…
Linked Issues check	✅ Passed	The changes satisfy issue #17908. They explicitly detach childless roots during clear(), centralize root-edge bookkeeping in detach_root_block(), preserve populated-tree removal behavior, and add CPU-…
Out of Scope Changes check	✅ Passed	The implementation and tests are limited to BlockRadixTree root cleanup, clear() termination, related lifecycle behavior, and regression coverage. No unrelated changes are described.

• Fix all pre-merge checks with AI

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[LiRunGuo]

Hi team, gentle follow-up on this PR — the fix for #17908 (BlockRadixTree.clear() hangs on childless RootBlock, blocking KVCacheManager.shutdown()) is ready for review. It is rebased on the latest main (917cbfc), CodeRabbit reported no actionable comments, and the new CPU-only regression tests pass locally alongside the existing kv_cache_manager_v2 unit tests, with a clean ruff run.

Could a maintainer trigger the standard CI pipeline (e.g. /bot run) when convenient, or let us know if anything else is needed from our side? Thanks!

[LiRunGuo]

Hi @lowsfer @eopXD, following up here with a direct ping, since you work on kv_cache_manager_v2.

This fixes #17908. BlockRadixTree.clear() spins forever when the tree holds a childless RootBlock, so KVCacheManager.shutdown() hangs. The change is limited to _block_radix_tree.py plus CPU-only unit tests, and DCO passes.

As an external contributor I can't start the CI pipeline myself. Could one of you trigger /bot run and take a look, or point me to the right reviewer? Thanks!

[nvpohanh]

[by Codex] @lowsfer Friendly reminder: could you review this PR? Thanks!

[LiRunGuo]

@nvpohanh thanks for routing this to @lowsfer last week. It's been quiet since, so if lowsfer is busy, is there someone else on the KV cache manager v2 side who could review it?

The change is still small (_block_radix_tree.py plus CPU-only unit tests), and the full CI pipeline hasn't been triggered yet, so it would also need someone to run /bot run. Thanks!

[nvpohanh]

[by Claude Code] @lowsfer Could you review this PR? Thanks!

[nvpohanh]

[by Claude Code] @lowsfer Could you review this PR? Thanks!

[karljang]

Thanks for the fix and your follow-ups, and sorry we left this waiting for review. The Python KVCacheManagerV2 implementation targeted by this PR was removed in #19154; the C++ implementation is now the supported path.

Would you be interested in checking whether the same edge case applies to C++ V2 and carrying over a regression test? Otherwise, we can close this PR as superseded by the removal.

[nvpohanh]

[by Claude Code] @lowsfer Could you review this PR? Thanks!

[nvpohanh]

[by Claude Code] @lowsfer Could you review this PR? Thanks!

[LiRunGuo]

Current main removed the Python backend of KVCacheManagerV2 in #19154 ([None][refactor] BREAKING: Remove the Python backend of KVCacheManagerV2). tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py and its tests no longer exist, so this fix has no rebase target.\n\nI also checked the C++ replacement at cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp: BlockRadixTree::clear() iterates the root map and clears it directly, so the childless-root non-termination from #17908 does not carry over. Closing this as obsolete rather than porting dead code. Please reopen or comment if you would like a C++ regression test for the equivalent semantics.