From 94ce0ab680a5eb9bf39cb458dc800a8824a80cda Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Fri, 31 Jul 2026 14:06:51 -0700 Subject: [PATCH 1/2] [nvbugs/6535767][fix] Size new radix-tree block from prev alone, not prev->prev addOrGetExistingBlock derived the new block's tokensPerBlock by calling prev->tokensPerBlock(), which reads prev->prev. NodeBase::detachNext sets prev = nullptr when eviction prunes a block that a live KvCache still holds via SharedPtr, so that grandparent hop dereferences null and segfaults in KvCache::_commitBlock. Derive the value from prev directly, mirroring the Python reference (_block_radix_tree.py:390-393), which computes the new block's own tokens_per_block one level up and never touches prev.prev outside an assert. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- .../kv_cache_manager_v2/blockRadixTree.cpp | 20 ++++++++++++++----- tests/integration/test_lists/waives.txt | 1 - 2 files changed, 15 insertions(+), 6 deletions(-) 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 5d0d4bc314c0..194357ba02e2 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 @@ -321,6 +321,17 @@ static bool isPrefix(std::vector const& prefix, std::vectorprev, which is null whenever `prev` has been detached +// from the tree by eviction while a live KvCache still holds it (see Block::prev in the header). +int childTokensPerBlock(NodeBase const& prev) noexcept +{ + if (prev.type() == NodeBase::Type::kROOT_BLOCK) + return prev.tokensPerBlock(); + return static_cast(static_cast(prev).tokens.size()); +} + } // anonymous namespace BlockKey Block::makeKey(BlockKey const& prevKey, TokenIdExt const* tokens, size_t count) @@ -343,10 +354,7 @@ Block::Block(BlockKey k, std::vector toks, NodeBase* prevNode, LifeC 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(static_cast(prev)->tokens.size()); + return childTokensPerBlock(*prev); } void Block::releasePages() @@ -520,7 +528,6 @@ SharedPtr 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 +539,9 @@ SharedPtr addOrGetExistingBlock( return it->second; } + // Not prev->tokensPerBlock() — see childTokensPerBlock. + int const tpb = childTokensPerBlock(*prev); + // 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) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 11da19460ee4..096683fa2142 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -91,7 +91,6 @@ accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_fp8_4gpus[attention_ accuracy/test_llm_api_pytorch.py::TestNemotronV3Super::test_nvfp4_parallelism[TP4_PP2] SKIP (https://nvbugs/6427411) accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_4gpus_static_eplb[moe_backend=CUTLASS] SKIP (https://nvbugs/6535767) accuracy/test_llm_api_pytorch.py::TestNemotronV3Ultra::test_nvfp4_parallelism[ADP2_PP2] SKIP (https://nvbugs/6427411) -accuracy/test_llm_api_pytorch.py::TestQwen3NextInstruct::test_nvfp4[tp1_block_reuse-cutlass] SKIP (https://nvbugs/6535767) accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_fp8[latency] SKIP (https://nvbugs/6177390) accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_fp8[throughput_latency] SKIP (https://nvbugs/6177390) accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_nvfp4[latency_moe_trtllm] SKIP (https://nvbugs/6402009) From 78ed9d7d8531ba989920e13832613c89645ebe93 Mon Sep 17 00:00:00 2001 From: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue, 4 Aug 2026 17:45:11 +0800 Subject: [PATCH 2/2] [nvbugs/6535767][fix] invalidate live KV reuse safely (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> --- .../kv_cache_manager_v2/blockRadixTree.cpp | 20 +- .../kv_cache_manager_v2/blockRadixTree.h | 3 +- .../kv_cache_manager_v2/kvCache.cpp | 51 +++-- .../kv_cache_manager_v2/kvCache.h | 4 + .../kv_cache_manager_v2/kvCacheManager.cpp | 4 + .../kv_cache_manager_v2/kvCacheManager.h | 2 +- .../kvCacheManagerV2StatsTest.cpp | 201 ++++++++++++++++++ 7 files changed, 256 insertions(+), 29 deletions(-) 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 194357ba02e2..c35d5fe57655 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 @@ -321,17 +321,6 @@ static bool isPrefix(std::vector const& prefix, std::vectorprev, which is null whenever `prev` has been detached -// from the tree by eviction while a live KvCache still holds it (see Block::prev in the header). -int childTokensPerBlock(NodeBase const& prev) noexcept -{ - if (prev.type() == NodeBase::Type::kROOT_BLOCK) - return prev.tokensPerBlock(); - return static_cast(static_cast(prev).tokens.size()); -} - } // anonymous namespace BlockKey Block::makeKey(BlockKey const& prevKey, TokenIdExt const* tokens, size_t count) @@ -347,14 +336,16 @@ Block::Block(BlockKey k, std::vector toks, NodeBase* prevNode, LifeC , tokens(std::move(toks)) , prev(prevNode) , storage(numLifeCycles, nullptr) + , mTokensPerBlock(prevNode->type() == NodeBase::Type::kROOT_BLOCK + ? prevNode->tokensPerBlock() + : static_cast(static_cast(prevNode)->tokens.size())) , mOrdinal(prevNode->ordinal() + 1) { } int Block::tokensPerBlock() const noexcept { - TLLM_CHECK_DEBUG_WITH_INFO(prev, "Block must have a prev"); - return childTokensPerBlock(*prev); + return mTokensPerBlock; } void Block::releasePages() @@ -539,8 +530,7 @@ SharedPtr addOrGetExistingBlock( return it->second; } - // Not prev->tokensPerBlock() — see childTokensPerBlock. - int const tpb = childTokensPerBlock(*prev); + 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. 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 c4a1debcb093..9292202bb161 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 @@ -136,7 +136,7 @@ struct NodeBase SharedPtr 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 void releasePages(); private: + int const mTokensPerBlock; BlockOrdinal mOrdinal; }; 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 6beb2bd4f9eb..fb62e6f65135 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 @@ -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,6 +1758,12 @@ void KvCache::commit(std::vector 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(mCommittedTokens.size()); + if (mHistoryLength < numCommitted) + setHistoryLength(numCommitted); + if (mCommitState == CommitState::VIRTUAL_STOP) { if (isEnd) @@ -1747,12 +1771,6 @@ void KvCache::commit(std::vector const& tokens, bool isEnd) 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(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. // --------------------------------------------------------------------------- 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 2fca81c8d17e..24a291d14d4d 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 @@ -441,6 +441,7 @@ class KvCache : public std::enable_shared_from_this std::optional id; // opaque identifier (mirrors Python's id field) private: + friend class KvCacheManager; friend class KvCacheIntrospection; friend std::vector batchedLockToGpu( KvCache& kvCache, std::vector const& targets); @@ -527,6 +528,9 @@ class KvCache : public std::enable_shared_from_this // 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. diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp index ae136b0306f7..9ef43fe23da4 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.cpp +++ b/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(); } diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h index a9737b5b282c..9e1cd96bf6d5 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h @@ -127,7 +127,7 @@ class KvCacheManager : public std::enable_shared_from_this 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 ------------------------------------------------- diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp index 2f7169638d29..cf0b980ed79f 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp @@ -17,6 +17,7 @@ #include "tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.h" #include "tensorrt_llm/batch_manager/kv_cache_manager_v2/config.h" +#include "tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.h" #include "tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCache.h" #include "tensorrt_llm/batch_manager/kv_cache_manager_v2/kvCacheManager.h" #include "tensorrt_llm/batch_manager/kv_cache_manager_v2/pendingStats.h" @@ -197,6 +198,206 @@ TEST(KvCacheManagerV2StatsTest, DisabledStatsSuppressManagerCommit) EXPECT_TRUE(manager->getAndResetIterationStats().empty()); } +TEST(KvCacheManagerV2LifecycleTest, DetachedBlockRetainsTokensPerBlock) +{ + auto const config = makeConfig(); + LifeCycleRegistry const lifeCycles{config}; + BlockRadixTree tree{lifeCycles, config.tokensPerBlock}; + RootBlock& root = tree.addOrGetExisting({}); + + auto block + = addOrGetExistingBlock(&root, lifeCycles.size(), std::vector(config.tokensPerBlock, TokenId{1})); + auto detached = root.detachNext(block->key); + ASSERT_EQ(detached, block); + EXPECT_TRUE(detached->isOrphan()); + EXPECT_TRUE(detached->isFull()); + EXPECT_EQ(detached->tokensPerBlock(), config.tokensPerBlock); + EXPECT_TRUE(detached->next.empty()); +} + +TEST(KvCacheManagerV2LifecycleTest, ResetPreventsLiveCacheFromRepopulatingReuse) +{ + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + auto manager = std::make_shared(makeConfig()); + CachedCudaStream stream; + auto cache = manager->createKvCache(); + ASSERT_TRUE(cache->resume(stream.handle())); + ASSERT_TRUE(cache->resize(4 * manager->tokensPerBlock())); + + std::vector const firstBlock(manager->tokensPerBlock(), TokenId{1}); + std::vector const secondBlock(manager->tokensPerBlock(), TokenId{2}); + std::vector const thirdBlock(manager->tokensPerBlock(), TokenId{3}); + std::vector const fourthBlock(manager->tokensPerBlock(), TokenId{4}); + std::vector committed = firstBlock; + committed.insert(committed.end(), secondBlock.begin(), secondBlock.end()); + cache->commit(committed); + EXPECT_EQ(manager->probeReuse({}, committed), 2 * manager->tokensPerBlock()); + auto const livePage = blockPageGetPage(cache->blocks()[BlockOrdinal{0}].pages[kDefaultBeamIndex][LifeCycleId{0}]); + auto const committedLivePage = dynamicPointerCast(livePage); + ASSERT_TRUE(committedLivePage); + ASSERT_TRUE(committedLivePage->hasValidSlot()); + + std::vector const siblingTokens(manager->tokensPerBlock(), TokenId{9}); + auto siblingCache = manager->createKvCache(); + ASSERT_TRUE(siblingCache->resume(stream.handle())); + ASSERT_TRUE(siblingCache->resize(manager->tokensPerBlock())); + siblingCache->commit(siblingTokens); + WeakPtr const siblingBlock = siblingCache->blocks()[BlockOrdinal{0}].treeBlock; + siblingCache->close(); + EXPECT_FALSE(siblingBlock.expired()); + EXPECT_EQ(manager->storage().getStatistics().evictable, SlotCount{1}); + + manager->clearReusableBlocks(); + EXPECT_EQ(cache->commitState(), KvCache::CommitState::VIRTUAL_STOP); + ASSERT_EQ(cache->blocks().size(), BlockOrdinal{4}); + EXPECT_TRUE(cache->blocks()[BlockOrdinal{0}].treeBlock->isOrphan()); + EXPECT_TRUE(cache->blocks()[BlockOrdinal{1}].treeBlock->isOrphan()); + EXPECT_TRUE(cache->blocks()[BlockOrdinal{0}].treeBlock->next.empty()); + EXPECT_TRUE(cache->blocks()[BlockOrdinal{1}].treeBlock->next.empty()); + EXPECT_EQ(cache->blocks()[BlockOrdinal{0}].treeBlock->getPage(LifeCycleId{0}), nullptr); + EXPECT_EQ(committedLivePage->block, nullptr); + EXPECT_EQ(committedLivePage->status(), PageStatus::LOCKED); + EXPECT_TRUE(committedLivePage->hasValidSlot()); + EXPECT_TRUE(siblingBlock.expired()); + EXPECT_EQ(manager->storage().getStatistics().evictable, SlotCount{0}); + EXPECT_TRUE(manager->radixTree().roots().empty()); + EXPECT_EQ(manager->probeReuse({}, committed), 0); + EXPECT_NO_THROW(cache->commit(thirdBlock)); + EXPECT_EQ(cache->numCommittedTokens(), 3 * manager->tokensPerBlock()); + EXPECT_EQ(cache->historyLength(), cache->numCommittedTokens()); + EXPECT_NO_THROW(cache->commit(fourthBlock)); + EXPECT_EQ(cache->numCommittedTokens(), 4 * manager->tokensPerBlock()); + EXPECT_EQ(cache->historyLength(), cache->numCommittedTokens()); + EXPECT_EQ(cache->commitState(), KvCache::CommitState::VIRTUAL_STOP); + EXPECT_TRUE(manager->radixTree().roots().empty()); + + cache->stopCommitting(); + EXPECT_EQ(cache->commitState(), KvCache::CommitState::USER_STOP); + cache->close(); + stream.synchronize(); + committed.insert(committed.end(), thirdBlock.begin(), thirdBlock.end()); + committed.insert(committed.end(), fourthBlock.begin(), fourthBlock.end()); + EXPECT_EQ(manager->probeReuse({}, committed), 0); +} + +TEST(KvCacheManagerV2LifecycleTest, OrphanedParentStopsFullBlockReuseBeforeV1Event) +{ + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + auto eventManager + = std::make_shared(1024, 0, std::nullopt, EventManager::AttentionDpGatherFn{}, "v1_block_key"); + auto manager = std::make_shared(makeConfig(), eventManager); + CachedCudaStream stream; + auto cache = manager->createKvCache(); + ASSERT_TRUE(cache->resume(stream.handle())); + ASSERT_TRUE(cache->resize(4 * manager->tokensPerBlock())); + + std::vector const firstBlock(manager->tokensPerBlock(), TokenId{1}); + cache->commit(firstBlock); + auto const orphan = cache->blocks()[BlockOrdinal{0}].treeBlock; + ASSERT_TRUE(orphan); + manager->radixTree().clear(); + ASSERT_TRUE(orphan->isOrphan()); + ASSERT_TRUE(manager->radixTree().roots().empty()); + + EXPECT_NO_THROW(cache->commit(std::vector(manager->tokensPerBlock(), TokenId{2}))); + EXPECT_EQ(cache->commitState(), KvCache::CommitState::VIRTUAL_STOP); + EXPECT_EQ(cache->numCommittedBlocks(), 1); + EXPECT_TRUE(orphan->next.empty()); + EXPECT_TRUE(manager->radixTree().roots().empty()); + EXPECT_NO_THROW(cache->commit(std::vector(manager->tokensPerBlock(), TokenId{3}))); + EXPECT_NO_THROW(cache->commit(std::vector(manager->tokensPerBlock(), TokenId{4}))); + EXPECT_EQ(cache->numCommittedTokens(), 4 * manager->tokensPerBlock()); + EXPECT_EQ(cache->historyLength(), cache->numCommittedTokens()); + + cache->stopCommitting(); + cache->close(); + stream.synchronize(); + EXPECT_NO_THROW(manager->radixTree().clear()); + EXPECT_NO_THROW(manager->shutdown()); +} + +TEST(KvCacheManagerV2LifecycleTest, OrphanedParentStopsPartialSnapshotBeforeV1Event) +{ + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + auto config = makeConfig(); + config.commitMinSnapshot = true; + auto eventManager + = std::make_shared(1024, 0, std::nullopt, EventManager::AttentionDpGatherFn{}, "v1_block_key"); + auto manager = std::make_shared(config, eventManager); + CachedCudaStream stream; + auto cache = manager->createKvCache(); + ASSERT_TRUE(cache->resume(stream.handle())); + ASSERT_TRUE(cache->resize(2 * manager->tokensPerBlock())); + + cache->commit(std::vector(manager->tokensPerBlock(), TokenId{1})); + auto const orphan = cache->blocks()[BlockOrdinal{0}].treeBlock; + ASSERT_TRUE(orphan); + manager->radixTree().clear(); + ASSERT_TRUE(orphan->isOrphan()); + + EXPECT_NO_THROW(cache->commit({TokenId{2}})); + EXPECT_EQ(cache->commitState(), KvCache::CommitState::VIRTUAL_STOP); + EXPECT_EQ(cache->numCommittedBlocks(), 1); + EXPECT_TRUE(orphan->next.empty()); + EXPECT_TRUE(manager->radixTree().roots().empty()); + + cache->close(); + stream.synchronize(); +} + +TEST(KvCacheManagerV2LifecycleTest, OrphanedParentStopsPartialFinalCommitBeforeV1Event) +{ + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + auto eventManager + = std::make_shared(1024, 0, std::nullopt, EventManager::AttentionDpGatherFn{}, "v1_block_key"); + auto manager = std::make_shared(makeConfig(), eventManager); + CachedCudaStream stream; + auto cache = manager->createKvCache(); + ASSERT_TRUE(cache->resume(stream.handle())); + ASSERT_TRUE(cache->resize(2 * manager->tokensPerBlock())); + + cache->commit(std::vector(manager->tokensPerBlock(), TokenId{1})); + auto const orphan = cache->blocks()[BlockOrdinal{0}].treeBlock; + ASSERT_TRUE(orphan); + manager->radixTree().clear(); + ASSERT_TRUE(orphan->isOrphan()); + + cache->commit({TokenId{2}}); + ASSERT_EQ(cache->commitState(), KvCache::CommitState::ALLOWED); + EXPECT_NO_THROW(cache->stopCommitting()); + EXPECT_EQ(cache->commitState(), KvCache::CommitState::USER_STOP); + EXPECT_EQ(cache->numCommittedBlocks(), 1); + EXPECT_TRUE(orphan->next.empty()); + EXPECT_TRUE(manager->radixTree().roots().empty()); + + cache->close(); + stream.synchronize(); +} + +TEST(KvCacheManagerV2LifecycleTest, ResetReleasesSwaPagesHeldOnlyForCommit) +{ + ASSERT_EQ(cudaSetDevice(0), cudaSuccess); + auto config = makeConfig(); + std::get(config.layers.front()).slidingWindowSize = 4; + auto manager = std::make_shared(config); + CachedCudaStream stream; + auto cache = manager->createKvCache(); + ASSERT_TRUE(cache->resume(stream.handle())); + ASSERT_TRUE(cache->resize(12, 0)); + ASSERT_TRUE(cache->resize(12, 8)); + + auto const& stalePage = cache->blocks()[BlockOrdinal{0}].pages[kDefaultBeamIndex][LifeCycleId{0}]; + ASSERT_TRUE(std::holds_alternative>(stalePage)); + manager->clearReusableBlocks(); + EXPECT_EQ(cache->commitState(), KvCache::CommitState::VIRTUAL_STOP); + EXPECT_TRUE(blockPageIsNull(stalePage)); + + EXPECT_NO_THROW(manager->clearReusableBlocks()); + EXPECT_TRUE(cache->resize(std::nullopt, 12)); + cache->close(); + stream.synchronize(); +} + TEST(KvCacheManagerV2StatsTest, PeakBlockStatsResetStartsNextIntervalFromCurrentSnapshot) { ASSERT_EQ(cudaSetDevice(0), cudaSuccess);