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..6663699b3ab3 --- /dev/null +++ b/tests/unittest/kv_cache_manager_v2_tests/test_nvbug_6625710.py @@ -0,0 +1,185 @@ +# 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 + + +@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.""" + + 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. + # + # 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), + "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 + # 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)