[https://nvbugs/6535767][fix] Safely invalidate live KV reuse

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp

@@ -336,17 +336,16 @@ Block::Block(BlockKey k, std::vector<TokenIdExt> toks, NodeBase* prevNode, LifeC
     , tokens(std::move(toks))
     , prev(prevNode)
     , storage(numLifeCycles, nullptr)
+    , mTokensPerBlock(prevNode->type() == NodeBase::Type::kROOT_BLOCK
+              ? prevNode->tokensPerBlock()
+              : static_cast<int>(static_cast<Block const*>(prevNode)->tokens.size()))
     , mOrdinal(prevNode->ordinal() + 1)
 {
 }
 
 int Block::tokensPerBlock() const noexcept
 {
-    TLLM_CHECK_DEBUG_WITH_INFO(prev, "Block must have a prev");
-    // Mirrors Python: prev.tokens_per_block if isinstance(prev, RootBlock) else len(prev.tokens)
-    if (prev->type() == Type::kROOT_BLOCK)
-        return prev->tokensPerBlock();
-    return static_cast<int>(static_cast<Block const*>(prev)->tokens.size());
+    return mTokensPerBlock;
 }
 
 void Block::releasePages()
@@ -520,7 +519,6 @@ SharedPtr<Block> addOrGetExistingBlock(
     }
 
     auto& prevNext = prev->next;
-    int const tpb = prev->tokensPerBlock();
     BlockKey newKey = Block::makeKey(prev->key, tokens.data(), tokens.size());
 
     // Exact match: return existing block (not new — mirrors Python's UselessBlockError path).
@@ -532,6 +530,8 @@ SharedPtr<Block> addOrGetExistingBlock(
         return it->second;
     }
 
+    int const tpb = prev->tokensPerBlock();
+
     // 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)

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h

@@ -136,7 +136,7 @@ struct NodeBase
 
     SharedPtr<Block> detachNext(BlockKey const& key);
 
-    /// RootBlock: delegates to tree. Block: len(prev->tokens) or prev->tokensPerBlock().
+    /// RootBlock: delegates to tree. Block: returns its retained block size.
     virtual int tokensPerBlock() const noexcept = 0;
 
 protected:
@@ -274,6 +274,7 @@ struct Block : NodeBase, EnableSharedFromThis<Block>
     void releasePages();
 
 private:
+    int const mTokensPerBlock;
     BlockOrdinal mOrdinal;
 };
 

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp

@@ -875,14 +875,19 @@ void KvCache::_snapshotPartialBlockToTree(BlockOrdinal ordinal, bool commitSsm)
 
     LifeCycleId numLc = mManager->storage().numLifeCycles();
     NodeBase* prevNode = nullptr;
-    RootBlock& root = mManager->radixTree().addOrGetExisting(mReuseScope);
     if (ordinal == BlockOrdinal{0})
     {
-        prevNode = &root;
+        prevNode = &mManager->radixTree().addOrGetExisting(mReuseScope);
     }
     else
     {
-        prevNode = _getTreeBlock(BlockOrdinal{ordinal.value() - 1}).get();
+        auto const& prevBlock = _getTreeBlock(BlockOrdinal{ordinal.value() - 1});
+        if (prevBlock->isOrphan())
+        {
+            _invalidateReuse();
+            return;
+        }
+        prevNode = prevBlock.get();
     }
 
     bool isNew = false;
@@ -1518,14 +1523,27 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm)
         throw LogicError("Cannot commit block that is not full except last block");
 
     // Prev node lookup (root or previous committed block).
-    RootBlock& root = mManager->radixTree().addOrGetExisting(mReuseScope);
     LifeCycleId numLc = mManager->storage().numLifeCycles();
 
-    NodeBase* prevNode = &root;
+    NodeBase* prevNode = nullptr;
     if (ord > 0)
     {
         TLLM_CHECK_DEBUG_WITH_INFO(mBlocks[BlockOrdinal{ord - 1}].treeBlock, "prev block must be committed");
-        prevNode = mBlocks[BlockOrdinal{ord - 1}].treeBlock.get();
+        auto const& prevBlock = mBlocks[BlockOrdinal{ord - 1}].treeBlock;
+        if (prevBlock->isOrphan())
+        {
+            _invalidateReuse();
+            if (isLast)
+            {
+                mCommitState = CommitState::USER_STOP;
+            }
+            return;
+        }
+        prevNode = prevBlock.get();
+    }
+    else
+    {
+        prevNode = &mManager->radixTree().addOrGetExisting(mReuseScope);
     }
 
     // Try to find or create a block in the radix tree.
@@ -1740,19 +1758,19 @@ void KvCache::commit(std::vector<TokenIdExt> const& tokens, bool isEnd)
 
     // Append tokens to committed list.
     mCommittedTokens.insert(mCommittedTokens.end(), tokens.begin(), tokens.end());
+
+    // Keep the history invariant even after reuse has been invalidated.
+    int const numCommitted = static_cast<int>(mCommittedTokens.size());
+    if (mHistoryLength < numCommitted)
+        setHistoryLength(numCommitted);
+
     if (mCommitState == CommitState::VIRTUAL_STOP)
     {
         if (isEnd)
             mCommitState = CommitState::USER_STOP;
         return;
     }
 
-    // Bump history_length to cover newly committed tokens (mirrors Python — done
-    // BEFORE the commit loop so stale-range computation sees the new history).
-    int const numCommitted = static_cast<int>(mCommittedTokens.size());
-    if (mHistoryLength < numCommitted)
-        setHistoryLength(numCommitted);
-
     int const numCommittedBlocksBefore = mNumCommittedBlocks;
     int const newNumFullBlocks = numCommitted / mTokensPerBlock;
     bool const hasPartialSnapshot = commitMinSnapshot && (numCommitted % mTokensPerBlock != 0) && numCommitted > 0;
@@ -1978,6 +1996,15 @@ void KvCache::_onStopCommitting()
     TLLM_CHECK_DEBUG(_checkSanity());
 }
 
+void KvCache::_invalidateReuse()
+{
+    if (mCommitState == CommitState::ALLOWED)
+    {
+        mCommitState = CommitState::VIRTUAL_STOP;
+        _onStopCommitting();
+    }
+}
+
 // ---------------------------------------------------------------------------
 // _setupForReuse: find existing blocks in radix tree matching input tokens.
 // ---------------------------------------------------------------------------

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h

@@ -441,6 +441,7 @@ class KvCache : public std::enable_shared_from_this<KvCache>
     std::optional<RequestIdType> id; // opaque identifier (mirrors Python's id field)
 
 private:
+    friend class KvCacheManager;
     friend class KvCacheIntrospection;
     friend std::vector<SharedPageLock> batchedLockToGpu(
         KvCache& kvCache, std::vector<BatchedLockTarget> const& targets);
@@ -527,6 +528,9 @@ class KvCache : public std::enable_shared_from_this<KvCache>
     // Mirrors Python's _on_stop_committing().
     void _onStopCommitting();
 
+    // Stop publishing reuse after a cache reset or detached parent invalidates this lineage.
+    void _invalidateReuse();
+
     // Commit a single block at ordinal `ord`.
     // `isLast` mirrors Python's is_last parameter: when True (or on VIRTUAL_STOP),
     // transitions to USER_STOP and calls _onStopCommitting() internally.

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp

@@ -151,6 +151,10 @@ void KvCacheManager::shutdown()
 void KvCacheManager::clearReusableBlocks()
 {
     TLLM_CHECK_DEBUG(mRadixTree);
+    for (KvCache* kvc : mLivingKvCaches)
+    {
+        kvc->_invalidateReuse();
+    }
     mRadixTree->clear();
 }
 

cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h

@@ -127,7 +127,7 @@ class KvCacheManager : public std::enable_shared_from_this<KvCacheManager>
 
     void shutdown();
 
-    // Clear all reusable (committed) blocks from the radix tree.
+    // Start a new reuse epoch: clear the tree and stop live caches from republishing into it.
     void clearReusableBlocks();
 
     // ---- KvCache creation -------------------------------------------------