[https://nvbugs/6625710][fix] Re-attach radix-tree blocks detached under a live request

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp

@@ -498,45 +498,53 @@ std::vector<SharedPtr<Block>> Block::clearStaleBlocksAfterPageUnlink(
     return detachedBlocks;
 }
 
-// ---------------------------------------------------------------------------
-// addOrGetExistingBlock
-// ---------------------------------------------------------------------------
-
-SharedPtr<Block> addOrGetExistingBlock(NodeBase* prev, std::vector<TokenIdExt> tokens, bool knownNoDigest, bool* isNew)
+SharedPtr<Block> getExistingBlock(NodeBase* prev, BlockKey const& key, TokenIdExt const* tokens, size_t numTokens)
 {
     TLLM_CHECK_DEBUG_WITH_INFO(prev, "prev must not be null");
 
-    // Prev must be a full block if it is a Block (mirrors Python: "prev must be a full block").
-    if (prev->type() == NodeBase::Type::kBLOCK)
-    {
-        TLLM_CHECK_DEBUG_WITH_INFO(static_cast<Block*>(prev)->isFull(), "prev must be a full block");
-    }
+    // Only a full block may be a parent; that is also why returning a covering sibling
+    // below is safe, since only partial blocks can be covered and never become parents.
+    TLLM_CHECK_DEBUG_WITH_INFO(
+        prev->type() != NodeBase::Type::kBLOCK || static_cast<Block*>(prev)->isFull(), "prev must be a full block");
 
     auto& prevNext = prev->next;
-    int const tpb = prev->tokensPerBlock();
-    BlockKey newKey = Block::makeKey(prev->key, tokens.data(), tokens.size(), knownNoDigest);
 
-    // Exact match: return existing block (not new — mirrors Python's UselessBlockError path).
-    auto it = prevNext.find(newKey);
+    // Exact match. On the re-attach path this is another request having re-committed the
+    // same prefix while we were detached; its key is identical, so the chain stays valid.
+    auto it = prevNext.find(key);
     if (it != prevNext.end())
     {
-        if (isNew)
-            *isNew = false;
         return it->second;
     }
 
-    // Useless check: is this block's token prefix covered by a sibling?
-    // Mirrors Python's UselessBlockError — throw with the sibling block.
-    if (static_cast<int>(tokens.size()) < tpb)
+    // Covered by a longer sibling: reuse it rather than insert a redundant shorter node.
+    // A short page on a longer block is well defined -- CommittedPage::numTokensInBlock
+    // records the span and canReplacePage() will not supersede a wider page.
+    if (static_cast<int>(numTokens) < prev->tokensPerBlock())
     {
         for (auto const& [k, sibling] : prevNext)
         {
-            if (sibling->tokens.size() >= tokens.size()
-                && isPrefix(tokens.data(), tokens.size(), sibling->tokens.data(), sibling->tokens.size()))
-                throw UselessBlockError(sibling);
+            if (sibling->tokens.size() >= numTokens
+                && isPrefix(tokens, numTokens, sibling->tokens.data(), sibling->tokens.size()))
+            {
+                return sibling;
+            }
         }
     }
 
+    return nullptr;
+}
+
+void attachBlock(NodeBase* prev, SharedPtr<Block> const& block)
+{
+    TLLM_CHECK_DEBUG_WITH_INFO(prev, "prev must not be null");
+    TLLM_CHECK_DEBUG(block);
+    // Precondition: nothing in the tree supersedes this block, so we cannot shadow a sibling.
+    TLLM_CHECK_DEBUG(getExistingBlock(prev, block->key, block->tokens.data(), block->tokens.size()) == nullptr);
+
+    auto& prevNext = prev->next;
+    auto const& tokens = block->tokens;
+
     // A later turn may extend a partial endpoint to this longer block, replacing the
     // partial sibling. That turn may not have a committable SWA page for this block:
     // commitMinSnapshot releases out-of-window pages, while SWA scratch reuse uses
@@ -557,13 +565,12 @@ SharedPtr<Block> addOrGetExistingBlock(NodeBase* prev, std::vector<TokenIdExt> t
     // would already have replaced the shorter one.
     TLLM_CHECK_DEBUG(toRemove.size() <= 1);
 
-    // Create the new block. ordinal, tokensPerBlock, and numLifeCycles are all
-    // derived from prev. Block stores the tokens as a plain vector (moved in).
-    auto block = makeShared<Block>(newKey, std::move(tokens), prev);
-
+    // Redundant for a freshly constructed block; the re-attach path needs it to restore
+    // the link the prune walk cleared.
+    block->prev = prev;
     // Keep the parent attached while covered children are replaced. Adding the replacement
     // first prevents detachNext() from pruning an emptied RootBlock out of the tree.
-    prevNext[newKey] = block;
+    prevNext[block->key] = block;
 
     for (auto const& k : toRemove)
     {
@@ -572,12 +579,56 @@ SharedPtr<Block> addOrGetExistingBlock(NodeBase* prev, std::vector<TokenIdExt> t
         block->adoptPagesFrom(*erasedBlock);
         TLLM_CHECK_DEBUG_WITH_INFO(erasedBlock->isOrphan(), "erased sibling must be orphan after removal");
     }
+}
 
+SharedPtr<Block> addOrGetExistingBlock(NodeBase* prev, std::vector<TokenIdExt> tokens, bool knownNoDigest, bool* isNew)
+{
+    TLLM_CHECK_DEBUG_WITH_INFO(prev, "prev must not be null");
+
+    BlockKey const newKey = Block::makeKey(prev->key, tokens.data(), tokens.size(), knownNoDigest);
+
+    // Query first so a block we would discard is never built. Must precede the move below.
+    if (auto existing = getExistingBlock(prev, newKey, tokens.data(), tokens.size()))
+    {
+        if (isNew)
+            *isNew = false;
+        return existing;
+    }
+
+    // ordinal, tokensPerBlock, and numLifeCycles are all derived from prev.
+    // Block stores the tokens as a plain vector (moved in).
+    auto block = makeShared<Block>(newKey, std::move(tokens), prev);
+    attachBlock(prev, block);
     if (isNew)
         *isNew = true;
     return block;
 }
 
+SharedPtr<Block> attachOrGetExistingBlock(NodeBase* prev, SharedPtr<Block> block, bool* attached)
+{
+    TLLM_CHECK_DEBUG(block);
+
+    if (auto existing = getExistingBlock(prev, block->key, block->tokens.data(), block->tokens.size()))
+    {
+        // Someone else installed an equivalent block while we held ours (on the re-attach
+        // path, another request re-committed this prefix during our orphan window). Hand
+        // over any pages the winner lacks rather than dropping them: ours are still valid
+        // for the same tokens, and adoptPagesFrom() keeps whichever page covers more.
+        if (existing != block && existing->ordinal() == block->ordinal())
+        {
+            existing->adoptPagesFrom(*block);
+        }
+        if (attached)
+            *attached = false;
+        return existing;
+    }
+
+    attachBlock(prev, block);
+    if (attached)
+        *attached = true;
+    return block;
+}
+
 // ---------------------------------------------------------------------------
 // removeSubtree
 // ---------------------------------------------------------------------------

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h

@@ -441,14 +441,29 @@ class BlockRadixTree
 // ---------------------------------------------------------------------------
 
 // Add a block to prev's `next` map, or return the existing one on collision.
-// Throws UselessBlockError (with the sibling block) if the block's tokens are a
-// prefix of an existing sibling — mirrors Python's UselessBlockError.
+// A partial block whose tokens are a prefix of an existing sibling returns that longer
+// sibling instead of inserting a redundant node.
 // If isNew is non-null, *isNew is set to true if a new block was created, false
 // if an existing block was returned.
 // knownNoDigest: from external text_only knowledge, never a scan (see Hasher::update).
 SharedPtr<Block> addOrGetExistingBlock(
     NodeBase* prev, std::vector<TokenIdExt> tokens, bool knownNoDigest, bool* isNew = nullptr);
 
+// Query: the block already in the tree that supersedes inserting `key`/`tokens` under
+// `prev` -- an exact-key match, or a longer sibling covering these tokens -- else nullptr.
+// Pure; lets callers avoid building a block they would discard.
+SharedPtr<Block> getExistingBlock(NodeBase* prev, BlockKey const& key, TokenIdExt const* tokens, size_t numTokens);
+
+// Mutation: link `block` under `prev` and absorb any covered shorter sibling.
+// Precondition (debug-asserted): getExistingBlock() returns nullptr for it.
+void attachBlock(NodeBase* prev, SharedPtr<Block> const& block);
+
+// The above two composed, for a caller that already holds a Block. Returns the block now
+// in the tree, which may be a pre-existing one; `attached` reports whether `block` itself
+// went in. Used by KvCache::_reattachOrphanTreeBlocks() to re-insert a block the tail-prune
+// walk detached while the request still held it -- see https://nvbugs/6625710.
+SharedPtr<Block> attachOrGetExistingBlock(NodeBase* prev, SharedPtr<Block> block, bool* attached = nullptr);
+
 // Post-order traversal: remove a subtree rooted at `root` from its parent's
 // next map. ~Block() handles page cleanup. Mirrors Python's remove_subtree().
 SharedPtr<Block> removeSubtree(Block& root);

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/exceptions.h

@@ -154,27 +154,6 @@ class OutOfPagesError : public std::runtime_error
     }
 };
 
-// Block creation rejected because its tokens are fully covered by an existing sibling.
-// Mirrors Python's UselessBlockError — carries the sibling block.
-// TODO: Once Python is removed and C++ becomes the primary development target,
-// replace this exception-based flow with a simple if-condition return in
-// addOrGetExistingBlock (returning the sibling block directly instead of throwing).
-// The exception pattern exists only to maintain parity with the Python code path.
-// Forward-declared; Block definition is in blockRadixTree.h.
-struct Block;
-
-class UselessBlockError : public std::runtime_error
-{
-public:
-    SharedPtr<Block> block;
-
-    explicit UselessBlockError(SharedPtr<Block> blk)
-        : std::runtime_error("Block is useless — covered by existing sibling")
-        , block(std::move(blk))
-    {
-    }
-};
-
 // ---------------------------------------------------------------------------
 // Helper: unwrap a weak_ptr, throw LogicError on dangling reference.
 // Mirrors Python's unwrap_rawref(_utils.py:163).

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp

@@ -846,6 +846,61 @@ void KvCache::_snapshotSsmToTreeBlock(SharedPtr<Block> const& treeBlock, LifeCyc
     _copyPageToTreeBlock(treeBlock, ssmLcId, srcPage, numTokensInBlock);
 }
 
+void KvCache::_reattachOrphanTreeBlocks(BlockOrdinal lastOrdinal, RootBlock& root)
+{
+    if (lastOrdinal < 0)
+    {
+        return;
+    }
+
+    // Every block at or below `lastOrdinal` is committed, so `treeBlock` is non-null:
+    // the walk below relies on that, since a null treeBlock would stop it early and then
+    // be dereferenced as `prevNode`.
+    auto isDetached = [this](BlockOrdinal ord)
+    {
+        auto const& sb = mBlocks[ord];
+        TLLM_CHECK_DEBUG_WITH_INFO(sb.treeBlock, "committed block must have a tree block");
+        return sb.treeBlock->isOrphan();
+    };
+
+    // Fast path: the block we are about to use as `prev` is still attached.
+    if (!isDetached(lastOrdinal))
+    {
+        return;
+    }
+
+    // Walk back to the deepest ancestor that is still in the tree (or the root).
+    BlockOrdinal first = lastOrdinal;
+    while (first >= 0 && isDetached(first))
+    {
+        --first;
+    }
+
+    NodeBase* prevNode
+        = (first < 0) ? static_cast<NodeBase*>(&root) : static_cast<NodeBase*>(mBlocks[first].treeBlock.get());
+
+    for (BlockOrdinal ord = first + 1; ord <= lastOrdinal; ++ord)
+    {
+        auto& sb = mBlocks[ord];
+
+        // detachNext() only cleared `prev` and the parent's map entry, so ordinal, tokens
+        // and surviving pages are intact -- re-attaching is enough, nothing to rebuild.
+        bool attached = false;
+        SharedPtr<Block> inTree = attachOrGetExistingBlock(prevNode, sb.treeBlock, &attached);
+        TLLM_CHECK_DEBUG(inTree);
+
+        if (attached && inTree->eventSink)
+        {
+            // Balances the addRemovedBlock() from detachNext(). If some other block was
+            // attached instead, it was already announced when it was created.
+            inTree->eventSink->addStoredBlock(*inTree);
+        }
+
+        sb.treeBlock = inTree;
+        prevNode = inTree.get();
+    }
+}
+
 // ---------------------------------------------------------------------------
 // _snapshotPartialBlockToTree: snapshot a partial final block into the radix
 // tree. Mirrors Python's _snapshot_partial_block_to_tree.
@@ -867,22 +922,18 @@ void KvCache::_snapshotPartialBlockToTree(BlockOrdinal ordinal, bool commitSsm)
     }
     else
     {
-        prevNode = _getTreeBlock(BlockOrdinal{ordinal.value() - 1}).get();
+        _reattachOrphanTreeBlocks(ordinal - 1, root);
+        prevNode = _getTreeBlock(ordinal - 1).get();
     }
 
     bool isNew = false;
-    SharedPtr<Block> treeBlock;
-    try
-    {
-        treeBlock = addOrGetExistingBlock(prevNode, tokens, textOnly(), &isNew);
-    }
-    catch (UselessBlockError const& e)
-    {
-        treeBlock = e.block;
-        isNew = false;
-    }
+    SharedPtr<Block> treeBlock = addOrGetExistingBlock(prevNode, tokens, textOnly(), &isNew);
     TLLM_CHECK_DEBUG(treeBlock);
-    TLLM_CHECK_DEBUG(isNew || std::equal(tokens.begin(), tokens.end(), treeBlock->tokens.begin()));
+    // When not new we either matched exactly or were given a longer sibling that covers
+    // these tokens, so our tokens are a prefix of the block we got back either way.
+    TLLM_CHECK_DEBUG(isNew
+        || (treeBlock->tokens.size() >= tokens.size()
+            && std::equal(tokens.begin(), tokens.end(), treeBlock->tokens.begin())));
 
     auto& beamBlock = mBlocks.at(ordinal).pages[kDefaultBeamIndex];
     auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId();
@@ -1510,27 +1561,20 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm)
     if (ord > 0)
     {
         TLLM_CHECK_DEBUG_WITH_INFO(mBlocks[BlockOrdinal{ord - 1}].treeBlock, "prev block must be committed");
+        _reattachOrphanTreeBlocks(BlockOrdinal{ord - 1}, root);
         prevNode = mBlocks[BlockOrdinal{ord - 1}].treeBlock.get();
     }
 
-    // Try to find or create a block in the radix tree.
-    // Mirrors Python's try/except UselessBlockError pattern.
-    // TODO: Replace with if-condition once Python is removed and C++ is the primary codebase.
+    // Find or create a block in the radix tree. A non-new result is either an exact match
+    // or a longer sibling that covers these tokens; both are usable here.
     bool blockIsNew = false;
-    SharedPtr<Block> newBlock;
-    try
-    {
-        newBlock = addOrGetExistingBlock(prevNode, tokenBlock, textOnly(), &blockIsNew);
-    }
-    catch (UselessBlockError const& e)
-    {
-        newBlock = e.block;
-        blockIsNew = false;
-    }
+    SharedPtr<Block> newBlock = addOrGetExistingBlock(prevNode, tokenBlock, textOnly(), &blockIsNew);
     TLLM_CHECK_DEBUG(newBlock);
     TLLM_CHECK_DEBUG(newBlock->tokensPerBlock() == mTokensPerBlock);
     // In reuse case, verify token match (mirrors Python: tree_block.tokens[:num_tokens] == tokens).
-    TLLM_CHECK_DEBUG(blockIsNew || std::equal(tokenBlock.begin(), tokenBlock.end(), newBlock->tokens.begin()));
+    TLLM_CHECK_DEBUG(blockIsNew
+        || (newBlock->tokens.size() >= tokenBlock.size()
+            && std::equal(tokenBlock.begin(), tokenBlock.end(), newBlock->tokens.begin())));
 
     auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId();
     bool didCommit = false;
@@ -2201,10 +2245,10 @@ bool KvCache::_checkSanity() const
                 }
                 else
                 {
-                    if (mStatus == Status::ACTIVE)
-                        TLLM_CHECK_DEBUG(std::holds_alternative<SharedPageLock>(bp));
-                    else
-                        TLLM_CHECK_DEBUG(std::holds_alternative<SharedPtr<PageHolder>>(bp));
+                    // The page must be present, and locked exactly when the cache is
+                    // active; a suspended cache holds it instead.
+                    TLLM_CHECK_DEBUG(!std::holds_alternative<std::monostate>(bp)
+                        && ((mStatus == Status::ACTIVE) == std::holds_alternative<SharedPageLock>(bp)));
                 }
 
                 if (!blockPageIsNull(bp))

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h

@@ -514,6 +514,13 @@ class KvCache : public std::enable_shared_from_this<KvCache>
     // no later writes to this KvCache's memory). Mirrors Python's _commit_block.
     void _commitBlock(int ord, bool isLast, bool commitSsm = false, bool moveSsm = false);
 
+    // Re-attach committed tree blocks of ours that the tail-prune walk detached while we
+    // still referenced them (triggered by evicting an unheld SSM snapshot). Relinks the
+    // chain from the deepest surviving ancestor. The evicted page stays absent, which is
+    // fine: pruneMatch() truncates reuse before a block that lacks it.
+    // See https://nvbugs/6625710.
+    void _reattachOrphanTreeBlocks(BlockOrdinal lastOrdinal, RootBlock& root);
+
     struct TakenPage
     {
         SharedPtr<UncommittedPage> page;

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp

@@ -1616,8 +1616,7 @@ void StorageManager::adjustCacheLevel(CacheLevel level, std::optional<size_t> ne
     }
     auto newNumSlots = lvlStorage.computeSlotCountList(ratioList, minSlots, quota);
 
-    if (!isLastLevel(level))
-        TLLM_CHECK_DEBUG(persistentPages == nullptr);
+    TLLM_CHECK_DEBUG(isLastLevel(level) || persistentPages == nullptr);
 
     // Shrink first.
     for (PoolGroupIndex pgIdx{0}; pgIdx < newNumSlots.size(); ++pgIdx)

cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp

@@ -1877,7 +1877,17 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m)
             {
                 parent = nb::cast<EventManagerTestBlock&>(parentObject).block.get();
             }
-            auto block = kv::addOrGetExistingBlock(parent, std::move(tokens), knownNoDigest);
+            // addOrGetExistingBlock() may hand back a pre-existing block -- an exact-key
+            // match, or a longer sibling covering these tokens. This helper then installs
+            // pages into `storage` directly (bypassing replacePage()), which would clobber
+            // that block's existing back-pointers, so reject the case outright: a test
+            // asking for a block that already exists is a test bug.
+            bool blockIsNew = false;
+            auto block = kv::addOrGetExistingBlock(parent, std::move(tokens), knownNoDigest, &blockIsNew);
+            if (!blockIsNew)
+            {
+                throw std::invalid_argument("make_test_block: an equivalent block is already in the tree");
+            }
 
             kv::TypedVec<kv::LifeCycleId, kv::SlotCount> counts(manager.lifeCycles().size(), 0);
             for (kv::LifeCycleId lifeCycle{0}; lifeCycle < manager.lifeCycles().size(); ++lifeCycle)