From 158a6298dc205adae0fca450108dd30df5fdc338 Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Sat, 22 Aug 2026 17:39:16 +0000 Subject: [PATCH 1/2] [https://nvbugs/6625710][fix] Re-attach radix-tree blocks detached under a live request An SSM snapshot is installed into the radix tree and immediately scheduleForEviction()-ed, so the eviction policy is its only owner -- no KvCache ever holds or locks it. The block it lives on is still referenced by the producing request through SeqBlock::treeBlock. Block::clearStaleBlocksAfterPageUnlink() prunes empty tail nodes. Evicting that unheld snapshot therefore detached a block a live request was still pointing at, and its next _commitBlock() dereferenced the now-null parent link and died in Block::tokensPerBlock(). The block is vulnerable only while it is still a leaf -- the request's own next commit gives it a child and closes the window -- which is why this surfaced as a ~1.7% flake. Requiring every life-cycle slot to be empty before pruning (already on main) covers hybrid attention+SSM models, where the attention page keeps a slot non-null. It does nothing for a pure-SSM model: with no attention life cycle the condition is trivially true and the block is detached exactly as before. KvCache::_reattachOrphanTreeBlocks() closes that gap. On commit, if the block we are about to use as `prev` was detached, walk back to the deepest surviving ancestor and re-attach the blocks we still hold. detachNext() only clears `prev` and the parent's map entry, so ordinal, tokens and surviving pages are all intact -- nothing has to be rebuilt. The evicted page stays absent, which is correct: pruneMatch() truncates reuse before a block that lacks it. Insertion is split so both paths share one implementation: getExistingBlock() -- pure query: the block already in the tree that supersedes this one, or nullptr. attachBlock() -- pure mutation: link under `prev` and absorb a covered shorter sibling. Debug-asserts the query is empty. addOrGetExistingBlock() (builds a block) and attachOrGetExistingBlock() (takes one) are thin wrappers over the pair, so the query now runs before construction and a block that would be discarded is never built. UselessBlockError is removed. A partial block covered by a longer sibling now returns that sibling instead of throwing; the two call sites that caught it did nothing but unwrap e.block, which is what the return value already expresses. The type never crossed the KVCM2 API boundary, so this is internal only. Tests: tests/unittest/kv_cache_manager_v2_tests/test_nvbug_6625710.py covers both configurations -- hybrid for the all-life-cycles condition, pure-SSM for the re-attach. pure_ssm segfaults without the fix and passes with it; it also asserts its own precondition, so retuning the quota or churn count cannot make it pass vacuously. The Python backend keeps its own copy of this logic and still needs both fixes. Signed-off-by: Yao Yao --- .../kv_cache_manager_v2/blockRadixTree.cpp | 105 ++++++++--- .../kv_cache_manager_v2/blockRadixTree.h | 19 +- .../kv_cache_manager_v2/exceptions.h | 21 --- .../kv_cache_manager_v2/kvCache.cpp | 104 ++++++++--- .../kv_cache_manager_v2/kvCache.h | 7 + .../kv_cache_manager_v2/storageManager.cpp | 3 +- .../batch_manager/kvCacheManagerV2.cpp | 12 +- .../test_nvbug_6625710.py | 171 ++++++++++++++++++ 8 files changed, 359 insertions(+), 83 deletions(-) create mode 100644 tests/unittest/kv_cache_manager_v2_tests/test_nvbug_6625710.py diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp index 4bd94a9e2383..863f9c285d25 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cpp @@ -498,45 +498,53 @@ std::vector> Block::clearStaleBlocksAfterPageUnlink( return detachedBlocks; } -// --------------------------------------------------------------------------- -// addOrGetExistingBlock -// --------------------------------------------------------------------------- - -SharedPtr addOrGetExistingBlock(NodeBase* prev, std::vector tokens, bool knownNoDigest, bool* isNew) +SharedPtr 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(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(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(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(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 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 addOrGetExistingBlock(NodeBase* prev, std::vector 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(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 addOrGetExistingBlock(NodeBase* prev, std::vector t block->adoptPagesFrom(*erasedBlock); TLLM_CHECK_DEBUG_WITH_INFO(erasedBlock->isOrphan(), "erased sibling must be orphan after removal"); } +} +SharedPtr addOrGetExistingBlock(NodeBase* prev, std::vector 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(newKey, std::move(tokens), prev); + attachBlock(prev, block); if (isNew) *isNew = true; return block; } +SharedPtr attachOrGetExistingBlock(NodeBase* prev, SharedPtr 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 // --------------------------------------------------------------------------- diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h index 0120d411ee36..7dcaa6ff69bd 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h +++ b/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 addOrGetExistingBlock( NodeBase* prev, std::vector 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 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 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 attachOrGetExistingBlock(NodeBase* prev, SharedPtr 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 removeSubtree(Block& root); diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/exceptions.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/exceptions.h index 428362e61af3..9b25954c6eea 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/exceptions.h +++ b/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; - - explicit UselessBlockError(SharedPtr 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). diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp index e29324712707..576ad5a5661e 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.cpp @@ -846,6 +846,61 @@ void KvCache::_snapshotSsmToTreeBlock(SharedPtr 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(&root) : static_cast(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 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 treeBlock; - try - { - treeBlock = addOrGetExistingBlock(prevNode, tokens, textOnly(), &isNew); - } - catch (UselessBlockError const& e) - { - treeBlock = e.block; - isNew = false; - } + SharedPtr 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 newBlock; - try - { - newBlock = addOrGetExistingBlock(prevNode, tokenBlock, textOnly(), &blockIsNew); - } - catch (UselessBlockError const& e) - { - newBlock = e.block; - blockIsNew = false; - } + SharedPtr 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(bp)); - else - TLLM_CHECK_DEBUG(std::holds_alternative>(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(bp) + && ((mStatus == Status::ACTIVE) == std::holds_alternative(bp))); } if (!blockPageIsNull(bp)) diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h index bd86f30fc6aa..dec45ba8aff1 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h @@ -514,6 +514,13 @@ class KvCache : public std::enable_shared_from_this // 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 page; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp index c99c0b2b164e..0ad542e60723 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/storageManager.cpp @@ -1616,8 +1616,7 @@ void StorageManager::adjustCacheLevel(CacheLevel level, std::optional 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) diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index a18f67cff29a..db667fc31fee 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -1877,7 +1877,17 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) { parent = nb::cast(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 counts(manager.lifeCycles().size(), 0); for (kv::LifeCycleId lifeCycle{0}; lifeCycle < manager.lifeCycles().size(); ++lifeCycle) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_nvbug_6625710.py b/tests/unittest/kv_cache_manager_v2_tests/test_nvbug_6625710.py new file mode 100644 index 000000000000..62e1d939055c --- /dev/null +++ b/tests/unittest/kv_cache_manager_v2_tests/test_nvbug_6625710.py @@ -0,0 +1,171 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Regression tests for https://nvbugs/6625710 (see also https://nvbugs/6553427). + +An SSM snapshot is installed into the radix tree and immediately +``scheduleForEviction()``-ed, so the eviction policy is its *only* owner -- no +``KvCache`` ever holds or locks it. The block it lives on is still referenced +by the producing request through ``SeqBlock::treeBlock``. + +``Block::clearStaleBlocksAfterPageUnlink()`` prunes empty tail nodes with + + while (curr && curr->next.empty() && curr->storage.at(lcIdx) == nullptr) + +so when the evicted page is the SSM snapshot the condition consults *only* the +SSM slot, and the block is detached even though a live request still references +it. That request is left holding a tree block whose ``prev`` is null; its next +``_commitBlock()`` dereferences the stale parent link and dies in +``Block::tokensPerBlock()``. + +The block is vulnerable only while it is still a *leaf* -- once the request +commits the following block ``next.empty()`` is false and the walk stops. In +production that window is short, hence the ~1.7% failure rate. These tests +hold it open by suspending the request; suspension is not part of the +mechanism, it only makes the race deterministic. + +Two variants, because the bug takes two fixes: + +* ``hybrid`` -- attention + SSM life cycles, as in Qwen3.5-35B-A3B, where the bug + was first seen. Covered by requiring *every* life-cycle slot to be empty before + pruning a tail node: the attention page keeps a slot non-null, so the walk stops + and the block is never detached. +* ``pure_ssm`` -- SSM life cycles only. With no attention life cycle that + all-life-cycles condition is trivially true and the block is detached exactly as + before, so this variant additionally requires re-attaching the blocks a live + request still holds. +""" + +import gc +import itertools +import unittest +from typing import cast + +# Import the module rather than `from ... import TestSSMSupport`: binding a TestCase +# subclass into this module's namespace makes pytest collect and re-run that entire +# sibling class here as well. +import test_kv_cache_manager_v2 as kv_test # type: ignore[import-not-found] +from test_kv_cache_manager_v2 import ( # type: ignore[import-not-found] + CachedCudaStream, + CudaStream, + GpuCacheTierConfig, + HostCacheTierConfig, + KVCacheManager, + TokenId, + TokenIdExt, + init_cuda_once, +) + +TOKENS_PER_BLOCK = 32 +PROMPT_BLOCKS = 2 +CHURN_REQUESTS = 100 + + +class TestNvBug6625710(unittest.TestCase): + """Evicting an unheld SSM snapshot must not detach a still-referenced block.""" + + def setUp(self) -> None: + init_cuda_once() + self._token_id_gen = itertools.count() + gc.collect() + gc.disable() + + def tearDown(self) -> None: + gc.enable() + if hasattr(self, "manager"): + del self.manager + + def next_token(self) -> TokenIdExt: + return TokenId(next(self._token_id_gen)) + + def _run(self, num_attn_layers: int, two_tier: bool = False) -> None: + # _make_ssm_config does not touch `self`; reuse it directly. + cfg = kv_test.TestSSMSupport._make_ssm_config( + self, + tokens_per_block=TOKENS_PER_BLOCK, + gpu_quota=4 << 20, + num_attn_layers=num_attn_layers, + num_ssm_layers=2, + ) + if two_tier: + # GPU pages migrate down to host, so the *host* tier is the last level -- + # that is where forceEvict() actually drops DROPPABLE pages. + cfg.cache_tiers = [ + GpuCacheTierConfig(quota=8 << 20), + HostCacheTierConfig(quota=8 << 20), + ] + self.manager = KVCacheManager(cfg) + stream_holder = CachedCudaStream() + stream = cast(CudaStream, stream_holder.handle) + + # --- Request A commits a prefix, snapshotting SSM state into the tree. + # From here the snapshot is owned solely by the eviction policy. + kv_a = self.manager.create_kv_cache() + self.assertTrue(kv_a.resume(stream)) + prompt: list[TokenIdExt] = [] + for i in range(PROMPT_BLOCKS): + kv_a.capacity = TOKENS_PER_BLOCK * (i + 1) + chunk = [self.next_token() for _ in range(TOKENS_PER_BLOCK)] + kv_a.commit(chunk) + prompt += chunk + + # A sleeps, holding its tree blocks. Its tail block is still a leaf. + kv_a.suspend() + + # --- "A lot of requests ran during its sleep": each commits snapshots of + # its own and closes, so the pool must reclaim A's unheld tail snapshot. + for _ in range(CHURN_REQUESTS): + kv = self.manager.create_kv_cache() + if not kv.resume(stream): + kv.close() + continue + for i in range(4): + kv.capacity = TOKENS_PER_BLOCK * (i + 1) + kv.commit([self.next_token() for _ in range(TOKENS_PER_BLOCK)]) + kv.close() + + # --- Precondition guard. The bug needs A's snapshot to have actually been + # evicted; if the churn above stopped triggering eviction (quota re-tuned, + # CHURN_REQUESTS lowered, eviction policy changed) the rest of this test + # would pass without exercising anything. With no attention life cycle, + # prefix reuse is only possible via an SSM snapshot, so a probe that can + # still reuse the whole prompt proves the snapshot survived. The hybrid + # config has no equivalent probe: its attention pages serve a prefix match + # whether or not the snapshot is gone. + if num_attn_layers == 0: + probe = self.manager.create_kv_cache(input_tokens=list(prompt)) + reusable = probe.num_committed_tokens + probe.close() + self.assertLess( + reusable, + len(prompt), + "A's SSM snapshot was not evicted, so this test is not exercising " + "https://nvbugs/6625710 -- retune gpu_quota/CHURN_REQUESTS", + ) + + # --- A wakes up (the pool is free again) and keeps generating on top of + # its committed prefix. If its tail block was detached, this commit walks + # a null `prev`. + self.assertTrue(kv_a.resume(stream)) + kv_a.capacity = len(prompt) + TOKENS_PER_BLOCK + kv_a.commit([self.next_token() for _ in range(TOKENS_PER_BLOCK)]) + + kv_a.close() + self.manager.shutdown() + + def test_hybrid_attention_and_ssm(self) -> None: + self._run(num_attn_layers=2, two_tier=True) + + def test_pure_ssm(self) -> None: + self._run(num_attn_layers=0) From 44fb0d241e990ecb9e6ceef4bd176cfd0fb5c38f Mon Sep 17 00:00:00 2001 From: Yao Yao Date: Mon, 24 Aug 2026 04:39:00 +0000 Subject: [PATCH 2/2] [https://nvbugs/6625710][test] Gate on the C++ backend and guard the hybrid precondition Two review follow-ups on the regression tests. Gate the class on the C++ backend. The Python backend keeps its own copy of the prune logic and is not fixed yet, so an explicit TLLM_KV_CACHE_MANAGER_V2_BACKEND=python run turned these into deterministic failures rather than skips. Guard the hybrid precondition. pruneMatch() truncates a reuse match at the last block carrying an SSM snapshot, so a probe that can no longer reuse the whole prompt proves the snapshot is gone -- that already guarded the pure-SSM variant. For hybrid the attention pages survive independently, so also assert the attention-only prefix (the length before hybrid pruning) still spans the prompt; without it a shortfall could just mean the attention pages were evicted too. An earlier comment claimed hybrid had no equivalent probe, which was wrong: _get_num_tokens_before_hybrid_pruning() is exposed on both backends. Signed-off-by: Yao Yao --- .../test_nvbug_6625710.py | 40 +++++++++++++------ 1 file changed, 27 insertions(+), 13 deletions(-) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_nvbug_6625710.py b/tests/unittest/kv_cache_manager_v2_tests/test_nvbug_6625710.py index 62e1d939055c..6663699b3ab3 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_nvbug_6625710.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_nvbug_6625710.py @@ -72,6 +72,10 @@ CHURN_REQUESTS = 100 +@unittest.skipUnless( + kv_test.KV_CACHE_MANAGER_V2_BACKEND == "cpp", + "the Python backend carries its own copy of this logic and is not fixed yet", +) class TestNvBug6625710(unittest.TestCase): """Evicting an unheld SSM snapshot must not detach a still-referenced block.""" @@ -138,20 +142,30 @@ def _run(self, num_attn_layers: int, two_tier: bool = False) -> None: # --- Precondition guard. The bug needs A's snapshot to have actually been # evicted; if the churn above stopped triggering eviction (quota re-tuned, # CHURN_REQUESTS lowered, eviction policy changed) the rest of this test - # would pass without exercising anything. With no attention life cycle, - # prefix reuse is only possible via an SSM snapshot, so a probe that can - # still reuse the whole prompt proves the snapshot survived. The hybrid - # config has no equivalent probe: its attention pages serve a prefix match - # whether or not the snapshot is gone. - if num_attn_layers == 0: - probe = self.manager.create_kv_cache(input_tokens=list(prompt)) - reusable = probe.num_committed_tokens - probe.close() - self.assertLess( - reusable, + # would pass without exercising anything. + # + # pruneMatch() truncates a reuse match at the last block carrying an SSM + # snapshot, so a probe that can no longer reuse the whole prompt proves the + # snapshot is gone. For hybrid, the attention pages survive independently, + # so also assert the attention-only prefix (the length *before* hybrid + # pruning) still spans the prompt -- otherwise a shortfall could just mean + # the attention pages were evicted too, which is not the case under test. + probe = self.manager.create_kv_cache(input_tokens=list(prompt)) + reusable = probe.num_committed_tokens + attn_only = probe._get_num_tokens_before_hybrid_pruning() + probe.close() + self.assertLess( + reusable, + len(prompt), + "A's SSM snapshot was not evicted, so this test is not exercising " + "https://nvbugs/6625710 -- retune gpu_quota/CHURN_REQUESTS", + ) + if num_attn_layers > 0: + self.assertEqual( + attn_only, len(prompt), - "A's SSM snapshot was not evicted, so this test is not exercising " - "https://nvbugs/6625710 -- retune gpu_quota/CHURN_REQUESTS", + "attention pages were evicted too, so the reuse shortfall above does " + "not isolate the missing SSM snapshot -- retune gpu_quota/CHURN_REQUESTS", ) # --- A wakes up (the pool is free again) and keeps generating on top of