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 212161acfda0..5d0d4bc314c0 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 @@ -393,6 +393,55 @@ int Block::partialMatchThisNode(TokenIdExt const* otherTokens, size_t otherCount return count; } +int Block::pageCoverage(LifeCycleId lcIdx) const +{ + auto const* page = getPage(lcIdx); + return page != nullptr ? page->numTokensInBlock : 0; +} + +bool Block::holdsPage(CommittedPage const& page) const +{ + return getPage(page.lifeCycle) == &page; +} + +bool Block::canReplacePage(LifeCycleId lcIdx, int numTokensInBlock) const +{ + auto const* existing = getPage(lcIdx); + return existing == nullptr || existing->numTokensInBlock < numTokensInBlock; +} + +void Block::replacePage(LifeCycleId lcIdx, CommittedPage* page) +{ + TLLM_CHECK_DEBUG(canReplacePage(lcIdx, page->numTokensInBlock)); + // Unlink first: excludeFromEviction() below may drop the eviction list's last + // reference and destroy the page, so the slot must already be empty. + auto* existing = unlinkPage(lcIdx); + if (existing != nullptr && existing->scheduledForEviction()) + { + existing->manager->excludeFromEviction(*existing); + existing = nullptr; // May be dangling now, set to nullptr + } + page->block = this; + storage.at(lcIdx) = page; +} + +void Block::adoptPagesFrom(Block& other) +{ + TLLM_CHECK_DEBUG(other.ordinal() == ordinal()); + for (LifeCycleId lcIdx{0}; lcIdx < storage.size(); ++lcIdx) + { + auto* page = other.getPage(lcIdx); + if (page == nullptr || !canReplacePage(lcIdx, page->numTokensInBlock)) + { + continue; + } + // Clear the source slot directly rather than via unlinkPage(), which would + // null the back-pointer replacePage() is about to overwrite. + other.storage.at(lcIdx) = nullptr; + replacePage(lcIdx, page); + } +} + CommittedPage* Block::unlinkPage(LifeCycleId lcIdx, CommittedPage* expectedPage) { auto& slot = storage.at(lcIdx); @@ -494,7 +543,12 @@ SharedPtr addOrGetExistingBlock( } } - // Remove siblings whose tokens are a strict prefix of ours. + // 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 + // temporary shared storage that is not preserved. Adopt the partial sibling's pages + // to keep the shorter endpoint reusable, retaining each page's recorded token count + // (see CommittedPage::numTokensInBlock). std::vector toRemove; for (auto const& [k, sibling] : prevNext) { @@ -504,18 +558,25 @@ SharedPtr addOrGetExistingBlock( toRemove.push_back(k); } } + // Two covered siblings would be prefixes of each other; the insertion logic + // would already have replaced the shorter one. + TLLM_CHECK_DEBUG(toRemove.size() <= 1); + + // Create the new block. ordinal and tokensPerBlock are derived from prev inside the Block ctor. + auto block = makeShared(newKey, std::move(tokens), prev, numLifeCycles); + + // 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; + for (auto const& k : toRemove) { auto erasedBlock = prev->detachNext(k); TLLM_CHECK_DEBUG(erasedBlock); + block->adoptPagesFrom(*erasedBlock); TLLM_CHECK_DEBUG_WITH_INFO(erasedBlock->isOrphan(), "erased sibling must be orphan after removal"); - (void) erasedBlock; } - // Create the new block. ordinal and tokensPerBlock are derived from prev inside the Block ctor. - auto block = makeShared(newKey, std::move(tokens), prev, numLifeCycles); - - prevNext[newKey] = block; if (isNew) *isNew = true; return block; @@ -664,11 +725,6 @@ int numMatchedTokens(std::vector const& matched, in return tokensPerBlock * (static_cast(matched.size()) - 1) + matched.back().numMatchedTokens; } -bool hasPage(Block const& block, LifeCycleId lcId) -{ - return block.storage.at(lcId) != nullptr; -} - } // anonymous namespace std::vector BlockRadixTree::matchTokenPath( @@ -732,70 +788,26 @@ std::vector BlockRadixTree::pruneMatch(std::vector< [this](auto const& m) { return m.numMatchedTokens == mTokensPerBlock; })); auto attnLcs = mLifeCycles.attentionLifeCycles(); - - // Full-attention layers require pages on every matched block. - std::vector fullAttnLcList; - for (auto [lcId, attn] : attnLcs) - { - if (!attn->windowSize.has_value()) - { - fullAttnLcList.push_back(lcId); - } - } - if (!fullAttnLcList.empty()) - { - int n = findIndex(matched.begin(), matched.end(), - [&](auto const& match) - { - return std::any_of(fullAttnLcList.begin(), fullAttnLcList.end(), - [&](LifeCycleId lcId) { return !hasPage(*match.block, lcId); }); - }); - matched.resize(static_cast(n)); - } - - std::vector> swaLcs; - for (auto [lcId, attn] : attnLcs) - { - if (attn->windowSize.has_value()) - { - swaLcs.push_back({lcId, attn}); - } - } - - // SWA sink blocks must all be available. - for (auto [lcId, attn] : swaLcs) - { - int const sinkBlocks = attn->numSinkBlocks; - int const limit = std::min(sinkBlocks, static_cast(matched.size())); - int n = findIndex(matched.begin(), matched.begin() + limit, - [&, lcId = lcId](auto const& match) { return !hasPage(*match.block, lcId); }); - if (n < sinkBlocks) - { - matched.resize(static_cast(n)); - } - } - auto ssmLcId = mLifeCycles.ssmLifeCycleId(); + + // Fixed-point loop: SSM may select an earlier exact snapshot, while attention may + // shorten the match to the coverage of a required page. Every retry strictly + // shortens the match, so the loop terminates. while (!matched.empty()) { + // Check SSM snapshot availability first: truncating to the last reusable SSM + // snapshot changes the matched length that all the attention checks use. if (ssmLcId.has_value()) { - // Truncate to the last block whose SSM snapshot is reusable at that - // block's matched-token count, then clamp the tail entry's matched - // token count to the snapshot length (mirrors _block_radix_tree.py). int ssmTrunc = 0; int ssmMatchLen = 0; for (int i = static_cast(matched.size()) - 1; i >= 0; --i) { - CommittedPage* page = matched[static_cast(i)].block->storage.at(*ssmLcId); - if (page == nullptr) - { - continue; - } - auto* ssmPage = dynamic_cast(page); - TLLM_CHECK_DEBUG(ssmPage != nullptr); - int const snapshotLen = ssmPage->numTokensInBlock; - if (matched[static_cast(i)].numMatchedTokens >= snapshotLen) + auto const& entry = matched[static_cast(i)]; + // An SSM page holds the recurrent state after exactly this many tokens, + // so reuse must stop there instead of anywhere inside the block. + int const snapshotLen = entry.block->pageCoverage(*ssmLcId); + if (snapshotLen > 0 && entry.numMatchedTokens >= snapshotLen) { ssmTrunc = i + 1; ssmMatchLen = snapshotLen; @@ -810,35 +822,49 @@ std::vector BlockRadixTree::pruneMatch(std::vector< matched.back().numMatchedTokens = ssmMatchLen; } + // Only pages that are active at this candidate endpoint constrain attention + // reuse. Full attention requires every block. SWA requires sink blocks and the + // trailing window, but not the stale blocks between them. In particular, at an + // exact block boundary with windowSize=1, every historical block is stale. int const numTok = numMatchedTokens(matched, mTokensPerBlock); - bool trimmed = false; - for (auto [lcId, attn] : swaLcs) + bool shortened = false; + for (auto [lcId, attn] : attnLcs) { - int n = findIndex(matched.rbegin(), matched.rend(), - [&, lcId = lcId](auto const& match) { return hasPage(*match.block, lcId); }); - if (n != 0) + auto const staleRange = attn->getStaleRange(numTok, mTokensPerBlock); + int const staleBeg = staleRange.beg.value(); + int const staleEnd = staleRange.end.value(); + int const numMatchedBlocks = static_cast(matched.size()); + for (int i = 0; i < numMatchedBlocks; ++i) { - matched.resize(matched.size() - static_cast(n)); - trimmed = true; + // Mirrors Python's chain(range(stale.beg), range(stale.end, len(matched))). + if (staleBeg <= i && i < staleEnd) + { + continue; + } + int const numMatched = matched[static_cast(i)].numMatchedTokens; + int const coverage = matched[static_cast(i)].block->pageCoverage(lcId); + if (coverage >= numMatched) + { + continue; + } + if (coverage > 0) + { + matched.resize(static_cast(i) + 1); + matched.back().numMatchedTokens = coverage; + } + else + { + matched.resize(static_cast(i)); + } + shortened = true; break; } - - auto staleRange = attn->getStaleRange(numTok, mTokensPerBlock); - BlockOrdinal const staleEnd = staleRange.end; - if (staleEnd < BlockOrdinal{static_cast(matched.size())}) + if (shortened) { - auto tailBegin = matched.begin() + static_cast(toSizeT(staleEnd)); - int nMissing = findIndex(matched.rbegin(), std::make_reverse_iterator(tailBegin), - [&, lcId = lcId](auto const& match) { return !hasPage(*match.block, lcId); }); - if (BlockOrdinal{static_cast(matched.size()) - nMissing} > staleEnd) - { - matched.resize(matched.size() - static_cast(nMissing) - 1); - trimmed = true; - break; - } + break; } } - if (!trimmed) + if (!shortened) { break; } 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 c3ae3035415c..c4a1debcb093 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 @@ -234,6 +234,31 @@ struct Block : NodeBase, EnableSharedFromThis return storage[lcIdx]; } + // Return the page's recorded token count, or zero if the slot is empty. For attention + // this is prefix coverage; for SSM it is an exact checkpoint position. + // Mirrors Python's Block.page_coverage(). + int pageCoverage(LifeCycleId lcIdx) const; + + // True when `page` currently occupies its lifecycle's slot in this block. + // Mirrors Python's Block.holds_page(). + bool holdsPage(CommittedPage const& page) const; + + // Whether a page recording `numTokensInBlock` may take over slot `lcIdx`. A slot + // keeps only the page with the largest recorded token count; for SSM that means only + // the latest checkpoint, and a rare second endpoint in one block is a reuse miss. + // Pure; use replacePage() to install. Mirrors Python's Block.can_replace_page(). + bool canReplacePage(LifeCycleId lcIdx, int numTokensInBlock) const; + + // Install `page` in slot `lcIdx`, detaching whatever it supersedes. The superseded + // page may outlive this call while a request still holds it, so unlinkPage() must + // clear its back-pointer — releasePages() walks `storage` and would never see it. + // Mirrors Python's Block.replace_page(). + void replacePage(LifeCycleId lcIdx, CommittedPage* page); + + // Move `other`'s pages into this block without changing their recorded token counts. + // Mirrors Python's Block._adopt_pages_from(). + void adoptPagesFrom(Block& other); + // Clear stale tree nodes after a lifecycle page has been unlinked. // Returns detached blocks that must stay alive until cleanup completes. static std::vector> clearStaleBlocksAfterPageUnlink( diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp index 73fec6f57cd1..301d941d1cfc 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/eventManager.cpp @@ -57,6 +57,14 @@ uint64_t hash64Mix(int64_t input, uint64_t seed) return seed ^ (value + static_cast(kHashCombineConst) + (seed << 6U) + (seed >> 2U)); } +// A page whose recorded token count is short of the block's span cannot be announced: the +// event payload carries the block's full token list and cannot express a shorter valid +// prefix. +bool pageCoversBlock(CommittedPage const* page, Block const& block) +{ + return page != nullptr && page->numTokensInBlock >= static_cast(block.tokens.size()); +} + } // namespace EventManager::EventManager(int maxKvEventEntries, int windowSize, std::optional attentionDpRank, @@ -176,7 +184,7 @@ void EventManager::addStoredBlockUnlocked(Block const& block) std::set lifeCycleIds; for (LifeCycleId lifeCycle{0}; lifeCycle < block.storage.size(); ++lifeCycle) { - if (block.storage[lifeCycle] != nullptr) + if (pageCoversBlock(block.getPage(lifeCycle), block)) { lifeCycleIds.insert(lifeCycle.value()); } @@ -538,8 +546,8 @@ std::optional EventManager::storedBlockFromBlock( { continue; } - auto const* page = block.storage[lifeCycle]; - if (page != nullptr) + auto const* page = block.getPage(lifeCycle); + if (pageCoversBlock(page, block)) { cacheLevel = page->cacheLevel; priority = page->priority; diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.cpp index 835c48c9a13b..4a167d3792f3 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.cpp @@ -21,6 +21,7 @@ #include "kv_cache_manager_v2/kvCache.h" #include "kv_cache_manager_v2/kvCacheManager.h" +#include "kv_cache_manager_v2/page.h" #include "kv_cache_manager_v2/storageManager.h" #include @@ -79,6 +80,21 @@ KvCacheIntrospection::ActivePageStats KvCacheIntrospection::activePageStats(KvCa return {std::move(counts), std::move(unscheduledEvictable)}; } +std::optional KvCacheIntrospection::committedPageIsLinked(KvCache const& kvCache, int ordinal, int lcId) +{ + auto page = kvCache._page(BlockOrdinal{ordinal}, kDefaultBeamIndex, LifeCycleId{lcId}); + if (!page) + { + return std::nullopt; + } + auto committed = dynamicPointerCast(page); + if (!committed) + { + return std::nullopt; + } + return committed->block != nullptr; +} + TypedVec KvCacheIntrospection::storageStatistics( KvCacheManager& manager, CacheLevel level) { diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.h index 07f0dd9acef0..3d11ab920221 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/introspection.h @@ -35,6 +35,11 @@ class KvCacheIntrospection using ActivePageStats = std::tuple, TypedVec>; static ActivePageStats activePageStats(KvCache const& kvCache); + + // Whether the sequence's page at (ordinal, lcId) still points at a tree block; + // nullopt when the slot is empty or holds an uncommitted page. Test hook for the + // back-pointer invariant Block::replacePage() maintains. + static std::optional committedPageIsLinked(KvCache const& kvCache, int ordinal, int lcId); static bool allTreePagesDroppable(KvCacheManager& manager); static TypedVec storageStatistics(KvCacheManager& manager, CacheLevel level); 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 1715f44c7463..6beb2bd4f9eb 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 @@ -751,6 +751,8 @@ void KvCache::_subtractPendingAllocationRange(BlockOrdinal blockBegin, BlockOrdi bool KvCache::_hasReuseSource(BlockPage const& page) { + // `block` is set only while the page occupies its block's slot, so this asks whether + // the tree still offers the page. auto const committedPage = dynamicPointerCast(blockPageGetPage(page)); return committedPage && committedPage->block != nullptr; } @@ -779,16 +781,13 @@ void KvCache::_clearBlocks() // attached to a radix tree block. Mirrors Python's _copy_page_to_tree_block. // --------------------------------------------------------------------------- -void KvCache::_copyPageToTreeBlock(SharedPtr const& treeBlock, LifeCycleId lcIdx, SharedPtr const& srcPage, - std::optional ssmNumTokensInBlock) +CommittedPage* KvCache::_copyPageToTreeBlock( + SharedPtr const& treeBlock, LifeCycleId lcIdx, SharedPtr const& srcPage, int numTokensInBlock) { - if (treeBlock->storage.at(lcIdx) != nullptr) + if (!treeBlock->canReplacePage(lcIdx, numTokensInBlock)) { - return; // block already holds a page for this lifecycle + return treeBlock->getPage(lcIdx); } - auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); - bool const isSsm = ssmLcId.has_value() && lcIdx == *ssmLcId; - TLLM_CHECK_DEBUG(isSsm == ssmNumTokensInBlock.has_value()); auto& storageMgr = mManager->storage(); PoolGroupIndex pgIdx = storageMgr.getPoolGroupIndex(lcIdx); @@ -810,25 +809,21 @@ void KvCache::_copyPageToTreeBlock(SharedPtr const& treeBlock, LifeCycleI newSlot.readyEvent.waitInStream(reinterpret_cast(stream)); copySlotData(storageMgr, lvl, srcPage->cacheLevel, pgIdx, newSlot.slotId(), srcPage->slotId(), stream); - CachedCudaEvent readyEv(reinterpret_cast(stream)); - auto tempPage = makeShared(*this, treeBlock->ordinal(), lcIdx, lvl, kDefaultBeamIndex); - tempPage->setSlot(newSlot); - SharedPtr committed; - if (ssmNumTokensInBlock.has_value()) - { - committed = tempPage->convertToSsmCommitted(treeBlock, std::move(readyEv), *ssmNumTokensInBlock); - } - else - { - committed = tempPage->convertToCommitted(treeBlock, std::move(readyEv)); - } + newSlot.readyEvent = CachedCudaEvent(reinterpret_cast(stream)); + auto committed = makeShared( + &storageMgr, treeBlock, lcIdx, lvl, numTokensInBlock, getPriority(treeBlock->ordinal(), lcIdx)); + committed->setSlot(newSlot); + // Drops the superseded page, deferred until the copy is issued: an + // OutOfPagesError above must not destroy a usable shorter snapshot. + treeBlock->replacePage(lcIdx, committed.get()); // Schedule for eviction so eviction controller keeps a strong reference, // preventing the page from being destroyed. storageMgr.scheduleForEviction(*committed); - return; // success + return committed.get(); // success } // No pages available in any level, silently skip snapshot (matches Python). + return nullptr; } // --------------------------------------------------------------------------- @@ -841,28 +836,10 @@ void KvCache::_snapshotSsmToTreeBlock(SharedPtr const& treeBlock, LifeCyc int const numTokensInBlock = numTokens - treeBlock->ordinal().value() * mTokensPerBlock; TLLM_CHECK_DEBUG(0 < numTokensInBlock && numTokensInBlock <= mTokensPerBlock); - CommittedPage* existingRaw = treeBlock->storage.at(ssmLcId); - int existingNumTokens = 0; - if (existingRaw != nullptr) - { - auto* existingSsm = dynamic_cast(existingRaw); - TLLM_CHECK_DEBUG(existingSsm != nullptr); - existingNumTokens = existingSsm->numTokensInBlock; - } - if (existingNumTokens >= numTokensInBlock) + if (treeBlock->pageCoverage(ssmLcId) >= numTokensInBlock) { return; } - if (existingRaw != nullptr) - { - // Detach the smaller snapshot; the CommittedPage dtor's expected-page guard - // makes the later unlink a no-op (the slot is already null). - treeBlock->storage.at(ssmLcId) = nullptr; - if (existingRaw->scheduledForEviction()) - { - existingRaw->manager->excludeFromEviction(*existingRaw); - } - } auto& ssmBlock = mSsmBlocks[kDefaultBeamIndex]; auto* ssmLock = std::get_if(&ssmBlock[ssmLcId]); @@ -874,7 +851,8 @@ void KvCache::_snapshotSsmToTreeBlock(SharedPtr const& treeBlock, LifeCyc auto up = dynamicPointerCast(unlocked); TLLM_CHECK_DEBUG(up != nullptr); ssmBlock[ssmLcId] = std::monostate{}; - auto committed = up->convertToSsmCommitted(treeBlock, finishEvent(), numTokensInBlock); + // convertToCommitted() installs via replacePage(), dropping any shorter snapshot. + auto committed = up->convertToCommitted(treeBlock, finishEvent(), numTokensInBlock); mManager->storage().scheduleForEviction(*committed); return; } @@ -923,21 +901,21 @@ void KvCache::_snapshotPartialBlockToTree(BlockOrdinal ordinal, bool commitSsm) auto& beamBlock = mBlocks.at(ordinal).pages[kDefaultBeamIndex]; auto ssmLcId = mManager->lifeCycles().ssmLifeCycleId(); - // Only attach partial attention pages to a tree block whose token span is - // exactly the partial snapshot. A longer existing sibling may already have - // full attention pages; otherwise a partial page would make it look more - // reusable than it is. - if (static_cast(treeBlock->tokens.size()) == numTokens) + // `treeBlock` may be a longer existing sibling that covers these tokens. Attaching the + // partial attention pages to it is still correct because each page records the token + // span it covers, and prefix matching honours that span. + std::vector attachedLcs; + for (auto const& [lcIdx, attn] : mManager->lifeCycles().attentionLifeCycles()) { - for (auto const& [lcIdx, attn] : mManager->lifeCycles().attentionLifeCycles()) + (void) attn; + auto& bp = beamBlock[lcIdx]; + if (blockPageIsNull(bp) || treeBlock->pageCoverage(lcIdx) >= numTokens) { - (void) attn; - auto& bp = beamBlock[lcIdx]; - if (blockPageIsNull(bp) || treeBlock->storage.at(lcIdx) != nullptr) - { - continue; - } - _copyPageToTreeBlock(treeBlock, lcIdx, blockPageGetPage(bp)); + continue; + } + if (_copyPageToTreeBlock(treeBlock, lcIdx, blockPageGetPage(bp), numTokens) != nullptr) + { + attachedLcs.push_back(lcIdx); } } if (commitSsm) @@ -945,9 +923,23 @@ void KvCache::_snapshotPartialBlockToTree(BlockOrdinal ordinal, bool commitSsm) TLLM_CHECK_DEBUG(ssmLcId.has_value()); _snapshotSsmToTreeBlock(treeBlock, *ssmLcId, start + numTokens); } - if (isNew && treeBlock->eventSink) + if (treeBlock->eventSink) { - treeBlock->eventSink->addStoredBlock(*treeBlock); + if (isNew) + { + treeBlock->eventSink->addStoredBlock(*treeBlock); + } + else + { + // The block was already announced, so report just the life cycles this snapshot + // added. The event manager itself drops the ones whose page does not span the + // whole block: the payload carries the block's full token list and cannot + // express a shorter valid prefix. + for (auto lcIdx : attachedLcs) + { + treeBlock->eventSink->addStoredLifeCycle(*treeBlock, lcIdx); + } + } } } @@ -1519,7 +1511,8 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm) int start = ord * mTokensPerBlock; int end = std::min(start + mTokensPerBlock, static_cast(mCommittedTokens.size())); std::vector tokenBlock(mCommittedTokens.begin() + start, mCommittedTokens.begin() + end); - bool isFull = static_cast(tokenBlock.size()) == mTokensPerBlock; + int const numTokens = static_cast(tokenBlock.size()); + bool const isFull = (numTokens == mTokensPerBlock); if (!isLast && !isFull) throw LogicError("Cannot commit block that is not full except last block"); @@ -1568,7 +1561,7 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm) auto& [up, locked] = taken[lc]; if (!up) continue; - auto committed = up->convertToCommitted(newBlock, finishEvent()); + auto committed = up->convertToCommitted(newBlock, finishEvent(), numTokens); if (locked) sb.pages[kDefaultBeamIndex][lc] = committed->lock(*this, kDefaultBeamIndex, static_cast(ord), lc); @@ -1595,7 +1588,14 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm) auto& bp = sb.pages[kDefaultBeamIndex][lc]; if (blockPageIsNull(bp)) continue; - auto* existingPage = newBlock->storage.at(lc); + // A page covering fewer tokens than this block spans (moved in from a shorter + // sibling) is NOT a substitute for our own full page: adopting it would feed + // uninitialized KV for the uncovered tail into a live request. + auto* existingPage = newBlock->getPage(lc); + if (existingPage != nullptr && existingPage->numTokensInBlock < numTokens) + { + existingPage = nullptr; + } bool isLocked = std::holds_alternative(bp); if (existingPage == nullptr) { @@ -1606,7 +1606,7 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm) if (up) { bp = std::monostate{}; - auto committed = up->convertToCommitted(newBlock, finishEvent()); + auto committed = up->convertToCommitted(newBlock, finishEvent(), numTokens); if (newBlock->eventSink) { newBlock->eventSink->addStoredLifeCycle(*newBlock, lc); @@ -1624,7 +1624,7 @@ void KvCache::_commitBlock(int ord, bool isLast, bool commitSsm, bool moveSsm) if (up) { bp = std::monostate{}; - auto committed = up->convertToCommitted(newBlock, finishEvent()); + auto committed = up->convertToCommitted(newBlock, finishEvent(), numTokens); if (newBlock->eventSink) { newBlock->eventSink->addStoredLifeCycle(*newBlock, lc); @@ -2126,8 +2126,10 @@ SharedPtr const& KvCache::_getTreeBlock(BlockOrdinal ordinal) const else if (!blockPageIsNull(beamBlock[lcId])) { auto page = blockPageGetPage(beamBlock[lcId]); + // committed->block may differ from `ret`: a page can be moved to a longer + // sibling block, or replaced there by one with larger token coverage. auto committed = dynamicPointerCast(page); - TLLM_CHECK(committed && committed->block == ret.get()); + TLLM_CHECK(committed); } } } @@ -2196,6 +2198,12 @@ bool KvCache::_checkSanity() const { TLLM_CHECK_DEBUG(blockPageIsNull(bp)); } + else if (_getScratchRange(lcs.getLifeCycle(lc), 0).contains(ordinal)) + { + // It is uncertain whether this block should hold a page: it may have + // been scratch-allocated by an earlier chunk, but the prior history + // length and capacity are not retained. + } else { // For the decoder-side disagg case, for the first step, we will skip the 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 bd090197433b..2fca81c8d17e 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 @@ -457,11 +457,11 @@ class KvCache : public std::enable_shared_from_this std::vector _getMatchedTokens(BlockRadixTree::ReuseMatch const& match) const; void _clearBlocks(); // Copy `srcPage` into a new committed page attached to `treeBlock` for lifecycle - // `lcIdx`. When `ssmNumTokensInBlock` is set, the copy is an SsmCommittedPage - // covering that many tokens; otherwise a plain attention CommittedPage. No-op if - // the block already holds a page for this lifecycle, or on OOM in all levels. - void _copyPageToTreeBlock(SharedPtr const& treeBlock, LifeCycleId lcIdx, SharedPtr const& srcPage, - std::optional ssmNumTokensInBlock = std::nullopt); + // `lcIdx`, recording `numTokensInBlock` as the page's token count. Returns the page + // now in the slot, or nullptr on OOM in all levels. No-op (returning the existing + // page) if the block already holds a page covering at least that many tokens. + CommittedPage* _copyPageToTreeBlock( + SharedPtr const& treeBlock, LifeCycleId lcIdx, SharedPtr const& srcPage, int numTokensInBlock); // Snapshot live SSM state to `treeBlock` reusable at `numTokens` committed tokens. // If `move`, the live SSM page is moved (not copied) into the tree — the caller diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp index 91f9921df948..01e34e589a03 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.cpp @@ -22,7 +22,6 @@ #include "kv_cache_manager_v2/storageManager.h" // for StorageManager #include "tensorrt_llm/common/assert.h" -#include namespace tensorrt_llm::batch_manager::kv_cache_manager_v2 { @@ -96,18 +95,13 @@ SharedPageLock Page::lock(KvCache& kvCache, BeamIndex beamIndex, BlockOrdinal or // CommittedPage // --------------------------------------------------------------------------- -CommittedPage::CommittedPage(StorageManager* mgr, SharedPtr blk, LifeCycleId lc, CacheLevel level, Priority prio) +CommittedPage::CommittedPage( + StorageManager* mgr, SharedPtr blk, LifeCycleId lc, CacheLevel level, int numTokensInBlock_, Priority prio) : Page(mgr, lc, level, prio) , block(blk.get()) -{ -} - -SsmCommittedPage::SsmCommittedPage( - StorageManager* mgr, SharedPtr blk, LifeCycleId lc, CacheLevel level, Priority prio, int numTokensInBlock_) - : CommittedPage(mgr, std::move(blk), lc, level, prio) , numTokensInBlock(numTokensInBlock_) { - TLLM_CHECK_DEBUG(numTokensInBlock_ > 0); + TLLM_CHECK_DEBUG(0 < numTokensInBlock_ && numTokensInBlock_ <= static_cast(blk->tokens.size())); } CommittedPage::~CommittedPage() @@ -172,17 +166,20 @@ UncommittedPage::~UncommittedPage() // Delegate slot release to Page::~Page(). } -SharedPtr UncommittedPage::convertToCommitted(SharedPtr blk, CachedCudaEvent readyEv) +SharedPtr UncommittedPage::convertToCommitted( + SharedPtr blk, CachedCudaEvent readyEv, int numTokensInBlock) { TLLM_CHECK_DEBUG(!scheduledForEviction()); - TLLM_CHECK_DEBUG_WITH_INFO( - blk->storage.at(lifeCycle) == nullptr, "Block slot for this lifecycle already has a committed page"); + // Check before building: replacePage() below drops the superseded page, so a failure + // in between must not lose a usable snapshot. + TLLM_CHECK_DEBUG_WITH_INFO(blk->canReplacePage(lifeCycle, numTokensInBlock), + "Block slot for this lifecycle already has a page covering more tokens"); TLLM_CHECK_DEBUG_WITH_INFO(status() == PageStatus::DROPPABLE, "Release holder/lock before converting"); // Set the ready event before transfer (matches Python: self.ready_event = ready_event). this->readyEvent = std::move(readyEv); - auto committed = makeShared(manager, blk, lifeCycle, cacheLevel, priority); + auto committed = makeShared(manager, blk, lifeCycle, cacheLevel, numTokensInBlock, priority); // Move slot id to the committed page; invalidate our slot. committed->setSlotId(slotId()); // asserts valid committed->readyEvent = std::move(readyEvent); @@ -193,31 +190,7 @@ SharedPtr UncommittedPage::convertToCommitted(SharedPtr bl TLLM_CHECK_DEBUG_WITH_INFO(committed->hasValidSlot(), "committed page must have a valid slot after transfer"); // Register in block storage. - blk->storage.at(lifeCycle) = committed.get(); - - return committed; -} - -SharedPtr UncommittedPage::convertToSsmCommitted( - SharedPtr blk, CachedCudaEvent readyEv, int numTokensInBlock) -{ - TLLM_CHECK_DEBUG(!scheduledForEviction()); - TLLM_CHECK_DEBUG_WITH_INFO( - blk->storage.at(lifeCycle) == nullptr, "Block slot for this lifecycle already has a committed page"); - TLLM_CHECK_DEBUG_WITH_INFO(status() == PageStatus::DROPPABLE, "Release holder/lock before converting"); - - this->readyEvent = std::move(readyEv); - - auto committed = makeShared(manager, blk, lifeCycle, cacheLevel, priority, numTokensInBlock); - committed->setSlotId(slotId()); // asserts valid - committed->readyEvent = std::move(readyEvent); - resetSlot(); - readyEvent = CachedCudaEvent::makeNull(); - - TLLM_CHECK_DEBUG(!hasValidSlot() && readyEvent.isClosed()); - TLLM_CHECK_DEBUG_WITH_INFO(committed->hasValidSlot(), "committed page must have a valid slot after transfer"); - - blk->storage.at(lifeCycle) = committed.get(); + blk->replacePage(lifeCycle, committed.get()); return committed; } @@ -244,13 +217,12 @@ PageHolder::~PageHolder() if (!page->scheduledForEviction()) manager->scheduleForEviction(*page); - // If the block is orphan, exclude from eviction immediately. - auto* cp = dynamic_cast(page.get()); - if (cp) - { - if (cp->block == nullptr || cp->block->isOrphan()) - manager->excludeFromEviction(*page); - } + // A page that no longer sits in its block's slot (orphaned block, or replaced by a + // page with a larger recorded token count) is unreachable for reuse, so keeping it + // in the eviction LRU would just pin a slot until memory pressure hits. + auto* cp = static_cast(page.get()); + if (cp->block == nullptr || cp->block->isOrphan() || !cp->block->holdsPage(*cp)) + manager->excludeFromEviction(*page); } else { diff --git a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h index 9f52dca1f0ee..fc3a02252471 100644 --- a/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h +++ b/cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/page.h @@ -95,11 +95,23 @@ class CommittedPage : public Page public: Block* block; + // Token count recorded for this page. It is usually block->tokens.size(), but a + // snapshot taken at an earlier token boundary may live in a block that spans more + // tokens — see addOrGetExistingBlock() and KvCache::_snapshotPartialBlockToTree(). + // + // Attention and SSM life cycles interpret it differently: + // * for attention pages, it is the number of leading tokens with valid per-token KV, + // so the page is reusable for any prefix up to that count (compare with `>=`); + // * for an SSM page, it is the exact recurrent-state checkpoint, so reuse must be + // truncated to exactly that boundary. + int numTokensInBlock; + // Number of outstanding PlannedDropHandles that intend to drop this page. // Mirrors Python's CommittedPage.planned_drop_count. int plannedDropCount{0}; - CommittedPage(StorageManager* mgr, SharedPtr blk, LifeCycleId lc, CacheLevel level, Priority prio); + CommittedPage(StorageManager* mgr, SharedPtr blk, LifeCycleId lc, CacheLevel level, int numTokensInBlock, + Priority prio); ~CommittedPage() override; @@ -109,22 +121,6 @@ class CommittedPage : public Page } }; -// --------------------------------------------------------------------------- -// SsmCommittedPage — a committed SSM snapshot page. -// -// Unlike attention CommittedPages (which always cover a full block), an SSM -// snapshot may cover only a prefix of its block. `numTokensInBlock` records how -// many tokens of the block this snapshot is reusable for. -// --------------------------------------------------------------------------- -class SsmCommittedPage : public CommittedPage -{ -public: - int numTokensInBlock; - - SsmCommittedPage(StorageManager* mgr, SharedPtr blk, LifeCycleId lc, CacheLevel level, Priority prio, - int numTokensInBlock); -}; - // --------------------------------------------------------------------------- // UncommittedPage — page associated with a live KvCache sequence. // --------------------------------------------------------------------------- @@ -147,11 +143,10 @@ class UncommittedPage : public Page // Convert this UncommittedPage into a CommittedPage and attach to `block`. // The UncommittedPage becomes invalid (slot transferred to CommittedPage). - SharedPtr convertToCommitted(SharedPtr block, CachedCudaEvent readyEvent); - - // Convert this UncommittedPage into an SsmCommittedPage covering - // `numTokensInBlock` tokens and attach to `block`. Invalidates this page. - SharedPtr convertToSsmCommitted( + // + // `numTokensInBlock` records the page's token count. See + // CommittedPage::numTokensInBlock for its attention and SSM interpretations. + SharedPtr convertToCommitted( SharedPtr block, CachedCudaEvent readyEvent, int numTokensInBlock); }; diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index 7cda4942c104..0ddce2a3a8f7 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -1521,6 +1521,8 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) return std::make_tuple(std::move(counts.raw()), std::move(unscheduledEvictable.raw())); }, nb::arg("kv_cache"), nb::call_guard()); + mIntrospection.def("committed_page_is_linked", &kv::KvCacheIntrospection::committedPageIsLinked, + nb::arg("kv_cache"), nb::arg("ordinal"), nb::arg("lc_id"), nb::call_guard()); mIntrospection.def("all_tree_pages_droppable", &kv::KvCacheIntrospection::allTreePagesDroppable, nb::arg("manager"), nb::call_guard()); mIntrospection.def( @@ -1609,8 +1611,7 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) }, nb::arg("manager"), nb::call_guard()); // Returns (num_tokens, pages) where pages[i] is None for a block with no page in - // this lifecycle, else (slot_id, num_tokens_in_block) with num_tokens_in_block = -1 - // for a non-SSM (attention) page. + // this lifecycle, else (slot_id, num_tokens_in_block). mIntrospection.def( "reuse_match_pages", [](kv::KvCacheManager& manager, nb::object reuseScope, nb::object tokens, int lcId, bool enablePartial) @@ -1627,17 +1628,13 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m) pages.reserve(matchResult.blocks.stdSize()); for (auto* block : matchResult.blocks) { - auto* page = block->storage.at(lc); + auto* page = block->getPage(lc); if (page == nullptr) { pages.emplace_back(std::nullopt); continue; } - int const slotId = page->slotId().value(); - int numTokensInBlock = -1; - if (auto* ssm = dynamic_cast(page)) - numTokensInBlock = ssm->numTokensInBlock; - pages.emplace_back(std::make_pair(slotId, numTokensInBlock)); + pages.emplace_back(std::make_pair(page->slotId().value(), page->numTokensInBlock)); } } return std::make_tuple(numTokens, std::move(pages)); diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp index cd91444e6028..2f7169638d29 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerV2StatsTest.cpp @@ -220,7 +220,8 @@ TEST(KvCacheManagerV2StatsTest, PeakBlockStatsResetStartsNextIntervalFromCurrent tokens.emplace_back(TokenId{token++}); } auto block = addOrGetExistingBlock(previous, LifeCycleId{1}, std::move(tokens)); - auto page = makeShared(&storage, block, lifeCycle, kGpuLevel, kPriorityDefault); + auto page = makeShared( + &storage, block, lifeCycle, kGpuLevel, static_cast(block->tokens.size()), kPriorityDefault); page->setSlot(slot); block->storage[lifeCycle] = page.get(); storage.scheduleForEviction(*page); @@ -311,7 +312,8 @@ TEST(KvCacheManagerV2StatsTest, MigrationAndLastTierDropRecordersReceiveExactPag tokens.emplace_back(TokenId{tokenBase++}); } auto block = addOrGetExistingBlock(previous, LifeCycleId{1}, std::move(tokens)); - auto page = makeShared(&storage, block, lifeCycle, kGpuLevel, kPriorityDefault); + auto page = makeShared( + &storage, block, lifeCycle, kGpuLevel, static_cast(block->tokens.size()), kPriorityDefault); page->setSlot(slot); block->storage[lifeCycle] = page.get(); storage.scheduleForEviction(*page); diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py index c47c591f0d06..66e225ee63ba 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py @@ -13,7 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import TYPE_CHECKING, Iterable, Iterator, NamedTuple, Sequence, TypeVar, cast +from itertools import chain +from typing import TYPE_CHECKING, Iterator, NamedTuple, Sequence, TypeVar, cast from . import rawref from ._cache_key import ( # noqa: F401 @@ -26,15 +27,7 @@ ) from ._common import NDEBUG, BlockOrdinal, PageStatus, TokenIdExt from ._life_cycle_registry import AttnLifeCycle, LifeCycle, LifeCycleId, LifeCycleRegistry -from ._utils import ( - TypedIndexList, - expect_type, - filled_list, - find_index, - map_optional, - typed_enumerate, - unwrap_rawref, -) +from ._utils import TypedIndexList, filled_list, map_optional, typed_range, unwrap_rawref if TYPE_CHECKING: from ._event_manager import KVCacheEventManager @@ -276,12 +269,20 @@ def __init__(self, tokens: Sequence[TokenIdExt], prev: "Block | RootBlock") -> N for b in prev.next.values(): if b.tokens[: len(tokens)] == tokens: raise UselessBlockError(b) - # If there are sibling blocks fully covered by this block, remove them. + # 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: + # commit_min_snapshot releases out-of-window pages, while SWA scratch reuse uses + # temporary shared storage that is not preserved. Adopt the partial sibling's + # pages to keep the shorter endpoint reusable, retaining each page's recorded token + # count (see CommittedPage.num_tokens_in_block). to_remove = [] for k, b in prev.next.items(): if len(b.tokens) < len(tokens) and tokens[: len(b.tokens)] == b.tokens: assert NDEBUG or (not b.is_full and b is not self and b.key == k and not b.next) to_remove.append(k) + # Two covered siblings would be prefixes of each other; the insertion logic + # would already have replaced the shorter one. + assert NDEBUG or len(to_remove) <= 1 event_manager = get_tree(prev).event_manager if to_remove else None # Keep RootBlock attached while covered children are replaced. Adding # the replacement first prevents detach_next() from pruning an @@ -290,11 +291,62 @@ def __init__(self, tokens: Sequence[TokenIdExt], prev: "Block | RootBlock") -> N for k in to_remove: b = detach_next(prev, k) assert isinstance(b, Block) + self._adopt_pages_from(b) if event_manager is not None: event_manager.add_removed_event(b.key) assert b.is_orphan # _KVCache may still hold it. # prev.next keeps a strong ref to this _Block, so no need to remove self from prev.next in __del__(). + def page_coverage(self, lc_idx: LifeCycleId) -> int: + """Return the page's recorded token count, or zero if the slot is empty. + + For attention this is prefix coverage; for SSM it is an exact checkpoint position. + """ + page = self.get_page(lc_idx) + return page.num_tokens_in_block if page is not None else 0 + + def holds_page(self, page: "CommittedPage") -> bool: + return self.get_page(page.life_cycle) is page + + def can_replace_page(self, lc_idx: LifeCycleId, num_tokens_in_block: int) -> bool: + """Whether a page recording `num_tokens_in_block` may take over slot `lc_idx`. + + A slot keeps only the page with the largest recorded token count. For attention, + greater coverage strictly dominates lesser coverage. For SSM, this deliberately + keeps only the latest checkpoint -- two conversation turns rarely end inside the + same block, and if they do, a reuse miss is acceptable. + + Pure; use replace_page() to install. + """ + existing = self.get_page(lc_idx) + return existing is None or existing.num_tokens_in_block < num_tokens_in_block + + def replace_page(self, lc_idx: LifeCycleId, page: "CommittedPage") -> None: + """Install `page` in slot `lc_idx`, detaching whatever it supersedes. + + The superseded page may outlive this call while a request still holds it, so + unlink_page() must clear its back-pointer: _release_pages() walks `storage`, so + nothing would clear it later and it would dangle once this block dies. + """ + assert NDEBUG or self.can_replace_page(lc_idx, page.num_tokens_in_block) + existing = self.unlink_page(lc_idx) + if existing is not None and existing.scheduled_for_eviction: + existing.manager.exclude_from_eviction(existing) + page.block = rawref.ref(self) + self.storage[lc_idx] = rawref.ref(page) + + def _adopt_pages_from(self, other: "Block") -> None: + """Move `other`'s pages into self without changing their recorded token counts.""" + assert other.ordinal == self.ordinal + for lc_idx in typed_range(self.num_life_cycles): + page = other.get_page(lc_idx) + if page is None or not self.can_replace_page(lc_idx, page.num_tokens_in_block): + continue + # Clear the source slot directly rather than via unlink_page(), which would + # null the back-pointer replace_page() is about to overwrite. + other.storage[lc_idx] = None + self.replace_page(lc_idx, page) + def _release_pages(self) -> None: """Reclaim every page held by this block. @@ -309,9 +361,9 @@ def _release_pages(self) -> None: ``Block`` alive past ``StorageManager`` teardown, after which ``page.manager`` would be a dangling reference. """ - for lc_idx, ref in typed_enumerate(self.storage): - if ref is not None and ref() is not None: - page = unwrap_rawref(ref) + for lc_idx in typed_range(self.num_life_cycles): + page = self.get_page(lc_idx) + if page is not None: self.unlink_page(lc_idx) if page.status == PageStatus.DROPPABLE: if page.scheduled_for_eviction: @@ -339,30 +391,39 @@ def prev(self) -> "Block | RootBlock": return unwrap_rawref(self._prev) def get_page(self, lc_idx: LifeCycleId) -> "CommittedPage | None": + """Return the page in slot `lc_idx`, or None when the slot is empty. + + A non-empty slot always resolves: CommittedPage.__del__ unlinks the page from its + block before invalidating its rawref, so `storage` never retains a dangling ref. + """ return map_optional(self.storage[lc_idx], lambda f: f()) def unlink_page( self, lc_idx: LifeCycleId, expected_page: "CommittedPage | None" = None - ) -> bool: - page_ref = self.storage[lc_idx] - if page_ref is None: - return False + ) -> "CommittedPage | None": + """Detach slot `lc_idx`, returning the page that was there, or None. + + The sole place a block-page link is severed. + """ + # Called from CommittedPage.__del__, which invalidates the page's rawref only + # afterwards, so the dying page is still reachable here. + page = self.get_page(lc_idx) + if page is None: + return None # Only unlink when the slot still holds the expected page. During rebase # another block with the same key may have replaced the stored page, and # unlinking then would clobber the newer page's back-pointer. - if expected_page is not None and page_ref() is not expected_page: - return False - page = page_ref() - if page is not None: - page.block = rawref.NULL + if expected_page is not None and page is not expected_page: + return None + page.block = rawref.NULL self.storage[lc_idx] = None - return True + return page @staticmethod def clear_stale_blocks_after_page_unlink( start: "Block", lc_idx: LifeCycleId, lc: LifeCycle ) -> None: - assert start.storage[lc_idx] is None + assert start.get_page(lc_idx) is None ordinal = start.ordinal tree = try_get_tree(start) event_manager = tree.event_manager if tree is not None else None @@ -375,7 +436,7 @@ def clear_stale_blocks_after_page_unlink( # But for simplicity, we leave it for now. curr = start while ( - (isinstance(curr, Block) and curr.storage[lc_idx] is None) + (isinstance(curr, Block) and curr.get_page(lc_idx) is None) and not curr.next and curr._prev() is not None ): @@ -470,14 +531,6 @@ def _num_matched_tokens(self, matched: list[tuple[Block, int]]) -> int: return 0 return self._tokens_per_block * (len(matched) - 1) + matched[-1][1] - @staticmethod - def _has_pages(block: Block, lc_list: Iterable[LifeCycleId]) -> bool: - return all(block.storage[lc] is not None for lc in lc_list) - - @staticmethod - def _has_page(block: Block, lc: LifeCycleId) -> bool: - return block.storage[lc] is not None - # yields tuples of (block, num_matched_tokens). num_matched_tokens should be equal to # tokens_per_block except the last one. def _match_token_path( @@ -512,49 +565,23 @@ def _prune_match(self, matched: list[tuple[Block, int]]) -> list[tuple[Block, in assert all(b[1] == tokens_per_block for b in matched[:-1]) life_cycles = self._life_cycles - - # check for full attention layers attn_life_cycles = list(life_cycles.attention_life_cycles()) - if any(lc.window_size is None for _, lc in attn_life_cycles): - lc_list = [lc_idx for lc_idx, lc in attn_life_cycles if lc.window_size is None] - - def check_no_pages(b: tuple[Block, int]) -> bool: - return not BlockRadixTree._has_pages(b[0], lc_list) - - n = find_index(matched, check_no_pages) - matched = matched[:n] - - swa_life_cycles = tuple( - (lc_idx, lc) for lc_idx, lc in attn_life_cycles if lc.window_size is not None - ) - # check for SWA sink - for lc_idx, lc in swa_life_cycles: - - def check_no_page_lc(b: tuple[Block, int]) -> bool: - return not BlockRadixTree._has_page(b[0], lc_idx) - - n = find_index(matched[: lc.num_sink_blocks], check_no_page_lc) - if n < lc.num_sink_blocks: - matched = matched[:n] - # Check SSM snapshot availability before SWA window constraints. - # Truncating to the last reusable SSM snapshot can change the matched - # length used by the SWA check. ssm_lc_id = life_cycles.ssm_life_cycle_id - if ssm_lc_id is not None: - from ._page import SsmCommittedPage + # Fixed-point loop: SSM may select an earlier exact snapshot, while attention may + # shorten the match to the coverage of a required page. Every retry strictly + # shortens the match, so the loop terminates. while matched: + # Check SSM snapshot availability first: truncating to the last reusable SSM + # snapshot changes the matched length that all the attention checks use. if ssm_lc_id is not None: ssm_trunc = 0 ssm_match_len = 0 for i in reversed(range(len(matched))): - block = matched[i][0] - page = map_optional(block.storage[ssm_lc_id], lambda f: f()) - if page is None: - continue - page = expect_type(SsmCommittedPage, page) - snapshot_len = page.num_tokens_in_block - if matched[i][1] >= snapshot_len: + # An SSM page holds the recurrent state after exactly this many tokens, + # so reuse must stop there instead of anywhere inside the block. + snapshot_len = matched[i][0].page_coverage(ssm_lc_id) + if snapshot_len > 0 and matched[i][1] >= snapshot_len: ssm_trunc = i + 1 ssm_match_len = snapshot_len break @@ -562,29 +589,30 @@ def check_no_page_lc(b: tuple[Block, int]) -> bool: if not matched: break matched[-1] = (matched[-1][0], ssm_match_len) - # SWA window check - num_tokens = self._num_matched_tokens(matched) - for lc_idx, lc in swa_life_cycles: - if lc.window_size is None: - continue - def check_has_page_lc(b: tuple[Block, int]) -> bool: - return BlockRadixTree._has_page(b[0], lc_idx) - - n = find_index(reversed(matched), check_has_page_lc) - if n != 0: - matched = matched[:-n] + # Only pages that are active at this candidate endpoint constrain attention + # reuse. Full attention requires every block. SWA requires sink blocks and the + # trailing window, but not the stale blocks between them. In particular, at an + # exact block boundary with window_size=1, every historical block is stale. + num_tokens = self._num_matched_tokens(matched) + shortened = False + for lc_idx, lc in attn_life_cycles: + stale = lc.get_stale_range(num_tokens, tokens_per_block) + for i in chain(range(stale.beg), range(stale.end, len(matched))): + block, num_matched = matched[i] + coverage = block.page_coverage(lc_idx) + if coverage >= num_matched: + continue + if coverage > 0: + matched = matched[: i + 1] + matched[-1] = (block, coverage) + else: + matched = matched[:i] + shortened = True break - _, stale_end = lc.get_stale_range(num_tokens, tokens_per_block) - - def has_no_page(b: tuple[Block, int]) -> bool: - return not BlockRadixTree._has_page(b[0], lc_idx) - - n = find_index(reversed(matched[stale_end:]), has_no_page) - if len(matched) - n > stale_end: - matched = matched[: len(matched) - n - 1] + if shortened: break - else: + if not shortened: break return matched diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py index da0500dc4c06..d71bf3db6acb 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py @@ -57,7 +57,6 @@ CommittedPage, Page, ScratchSlotLock, - SsmCommittedPage, UncommittedPage, _PageHolder, _SharedPageLock, @@ -75,7 +74,6 @@ filled_list, intersect, make_typed, - map_optional, stream_wait_events, to_typed, typed_enumerate, @@ -83,7 +81,6 @@ typed_map, typed_range, unwrap_optional, - unwrap_rawref, value_or, ) from ._moving_average import Average @@ -1441,13 +1438,10 @@ def _copy_page_to_tree_block( tree_block: Block, lc_idx: LifeCycleId, src_page: Page, - ssm_num_tokens_in_block: int | None = None, + num_tokens_in_block: int, ) -> CommittedPage | None: - existing_page = map_optional(tree_block.storage[lc_idx], lambda p: p()) - if existing_page is not None: - return existing_page - is_ssm = lc_idx == self.manager._life_cycles.ssm_life_cycle_id - assert is_ssm == (ssm_num_tokens_in_block is not None) + if not tree_block.can_replace_page(lc_idx, num_tokens_in_block): + return tree_block.get_page(lc_idx) storage = self.manager._storage pg_idx = storage.get_pool_group_index(lc_idx) @@ -1471,19 +1465,12 @@ def _copy_page_to_tree_block( ) new_slot.ready_event = CachedCudaEvent(cuda_stream) priority = self._get_priority(tree_block.ordinal, self.manager._life_cycles[lc_idx]) - if ssm_num_tokens_in_block is None: - committed = CommittedPage(storage, tree_block, lc_idx, lvl, new_slot, priority) - else: - committed = SsmCommittedPage( - storage, - tree_block, - lc_idx, - lvl, - new_slot, - priority, - ssm_num_tokens_in_block, - ) - tree_block.storage[lc_idx] = rawref.ref(committed) + committed = CommittedPage( + storage, tree_block, lc_idx, lvl, new_slot, num_tokens_in_block, priority + ) + # Drops the superseded page, deferred until the copy is issued: an + # OutOfPagesError above must not destroy a usable shorter snapshot. + tree_block.replace_page(lc_idx, committed) storage.schedule_for_eviction(committed) return committed return None @@ -1495,17 +1482,8 @@ def _snapshot_ssm_to_tree_block( tokens_per_block = self.tokens_per_block num_tokens_in_block = num_tokens - tree_block.ordinal * tokens_per_block assert 0 < num_tokens_in_block <= tokens_per_block - existing_page = map_optional(tree_block.storage[ssm_lc_id], lambda p: p()) - existing_num_tokens = 0 - if existing_page is not None: - existing_ssm_page = expect_type(SsmCommittedPage, existing_page) - existing_num_tokens = existing_ssm_page.num_tokens_in_block - if existing_num_tokens >= num_tokens_in_block: + if tree_block.page_coverage(ssm_lc_id) >= num_tokens_in_block: return - if existing_page is not None: - tree_block.storage[ssm_lc_id] = None - if existing_page.scheduled_for_eviction: - existing_page.manager.exclude_from_eviction(existing_page) ssm_block = self._ssm_blocks[DEFAULT_BEAM_INDEX] ssm_lock = expect_type(_SharedPageLock, ssm_block[ssm_lc_id]) @@ -1513,19 +1491,15 @@ def _snapshot_ssm_to_tree_block( if move: src_page = expect_type(UncommittedPage, ssm_lock.unlock()) ssm_block[ssm_lc_id] = None - committed = src_page.convert_to_ssm_committed( + # convert_to_committed() reserves the slot, dropping any shorter snapshot. + committed = src_page.convert_to_committed( tree_block, self.finish_event, num_tokens_in_block ) storage = self.manager._storage storage.schedule_for_eviction(committed) return - self._copy_page_to_tree_block( - tree_block, - ssm_lc_id, - src_page, - ssm_num_tokens_in_block=num_tokens_in_block, - ) + self._copy_page_to_tree_block(tree_block, ssm_lc_id, src_page, num_tokens_in_block) def _snapshot_partial_block_to_tree(self, ordinal: BlockOrdinal, commit_ssm: bool) -> None: tokens_per_block = self.tokens_per_block @@ -1549,23 +1523,33 @@ def _snapshot_partial_block_to_tree(self, ordinal: BlockOrdinal, commit_ssm: boo beam_idx = DEFAULT_BEAM_INDEX beam_block = self._blocks[ordinal].pages[beam_idx] ssm_lc_id = self.manager._life_cycles.ssm_life_cycle_id - # Only attach partial attention pages to a tree block whose token span is - # exactly the partial snapshot. A longer existing sibling may already - # have full attention pages; if not, a partial attention page would make - # that longer block look more reusable than it is. - if len(tree_block.tokens) == num_tokens: - for lc_idx, _ in self.manager._life_cycles.attention_life_cycles(): - holder = beam_block[lc_idx] - if holder is None or tree_block.storage[lc_idx] is not None: - continue - self._copy_page_to_tree_block(tree_block, lc_idx, holder.page) + # `tree_block` may be a longer existing sibling that covers these tokens. Attaching + # the partial attention pages to it is still correct because each page records the + # token span it covers, and prefix matching honours that span. + attached_lcs = list[LifeCycleId]() + for lc_idx, _ in self.manager._life_cycles.attention_life_cycles(): + holder = beam_block[lc_idx] + if holder is None or tree_block.page_coverage(lc_idx) >= num_tokens: + continue + if ( + self._copy_page_to_tree_block(tree_block, lc_idx, holder.page, num_tokens) + is not None + ): + attached_lcs.append(lc_idx) if commit_ssm: assert ssm_lc_id is not None self._snapshot_ssm_to_tree_block(tree_block, ssm_lc_id, start + num_tokens) - if is_new: - event_manager = self.manager.event_manager - if event_manager is not None: + event_manager = self.manager.event_manager + if event_manager is not None: + if is_new: event_manager.add_stored_block_event_from_block(tree_block) + else: + # The block was already announced, so report just the life cycles this + # snapshot added. The event manager itself drops the ones whose page does + # not span the whole block: the payload carries the block's full token + # list and cannot express a shorter valid prefix. + for lc_idx in attached_lcs: + event_manager.add_stored_life_cycle_event_from_block(tree_block, int(lc_idx)) def _commit_block( self, @@ -1626,8 +1610,7 @@ def _commit_block( for lc, (page, locked) in typed_enumerate(uncommitted_pages): if page is None: continue - p = page.convert_to_committed(tree_block, self.finish_event) - tree_block.storage[lc] = rawref.ref(p) + p = page.convert_to_committed(tree_block, self.finish_event, num_tokens) # The page comes from uncommitted page of self, so safe to skip wait. beam_block[lc] = ( p.lock(self, beam_idx, ordinal, lc, skip_wait=True) if locked else p.hold() @@ -1648,13 +1631,24 @@ def _commit_block( continue # SSM pages are not rebased if beam_block[lc] is None: continue - existing_page = map_optional(tree_block.storage[lc], lambda p: p()) + # A page covering fewer tokens than this block spans (moved in from a + # shorter sibling) is NOT a substitute for our own full page: adopting it + # would feed uninitialized KV for the uncovered tail into a live request. + existing_page = tree_block.get_page(lc) + if existing_page is not None and existing_page.num_tokens_in_block < num_tokens: + existing_page = None locked = isinstance(beam_block[lc], _SharedPageLock) if existing_page is None: # The reusable page is gone. We put our own page into the tree block. - page = cast(UncommittedPage, cast(_SharedPageLock, beam_block[lc]).page) + # Keep this a single expression: a local holding the lock/holder would + # keep it alive past `beam_block[lc] = None`, and convert_to_committed() + # requires the page to be droppable by then. + page = cast( + UncommittedPage, + cast("_SharedPageLock | _PageHolder", beam_block[lc]).page, + ) beam_block[lc] = None - p = page.convert_to_committed(tree_block, self.finish_event) + p = page.convert_to_committed(tree_block, self.finish_event, num_tokens) event_manager = self.manager.event_manager if event_manager is not None: event_manager.add_stored_life_cycle_event_from_block(tree_block, int(lc)) @@ -1840,7 +1834,9 @@ def _get_tree_block(self, ordinal: BlockOrdinal) -> Block: if lc == ssm_lc_id: assert b is None # SSM pages live in _ssm_blocks elif b is not None: - assert isinstance(b.page, CommittedPage) and b.page.block() is ret + # b.page.block() may differ from `ret`: a page can be moved to a longer + # sibling block, or replaced there by one with larger token coverage. + assert isinstance(b.page, CommittedPage) return ret def _take_uncommitted_page( @@ -1896,9 +1892,15 @@ def get_range(lc: LifeCycle): assert holder is None continue start, end = stale_ranges[lc] + lc_obj = self.manager._life_cycles[lc] if start <= ordinal < end: if is_committed or self._commit_state != self.CommitState.ALLOWED: assert holder is None + elif ordinal in self._get_scratch_range(lc_obj, 0): + # It is uncertain whether this block should hold a page: it may + # have been scratch-allocated by an earlier chunk, but the prior + # history length and capacity are not retained. + pass else: # For the decoder-side disagg case, for the first step, we will skip the # out-of-window blocks. @@ -1908,7 +1910,6 @@ def get_range(lc: LifeCycle): ) else: # Scratch blocks have None pages but valid base_page_indices - lc_obj = self.manager._life_cycles[lc] sr = self._get_scratch_range(lc_obj) is_scratch = ordinal in sr if is_scratch: @@ -2044,7 +2045,7 @@ def _setup_for_reuse(self, match: ReuseMatch) -> None: typed_range(stale_start), typed_range(stale_end, BlockOrdinal(len(matched))) ): block = self._block(ordinal, beam_idx) - holder = unwrap_rawref(unwrap_optional(matched[ordinal].storage[lc_idx])).hold() + holder = unwrap_optional(matched[ordinal].get_page(lc_idx)).hold() # For partial blocks (last block, not full), we defer the copy to first resume(). # Just store the holder of the original committed page for now. block[lc_idx] = holder @@ -2068,11 +2069,11 @@ def _setup_for_reuse(self, match: ReuseMatch) -> None: # SSM reuse: hold the snapshot from the last matched block. Copy is deferred to first resume(). if ssm_lc_id is not None and matched: snapshot_block = matched[-1] - snapshot_ref = snapshot_block.storage[ssm_lc_id] - assert snapshot_ref is not None, ( + snapshot_page = snapshot_block.get_page(ssm_lc_id) + assert snapshot_page is not None, ( "Last matched block must have SSM snapshot after truncation" ) - snapshot_holder = unwrap_rawref(snapshot_ref).hold() + snapshot_holder = snapshot_page.hold() self._ssm_blocks[DEFAULT_BEAM_INDEX][ssm_lc_id] = snapshot_holder if should_record_stats and ssm_lc_id is not None: changed = self._pending_stats.record_ssm_snapshot_lookup( diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py index ad1cb3c12012..c493918ca367 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py @@ -335,7 +335,7 @@ def get_latest_events(self, timeout_ms: float | None = None) -> list[KVCacheEven def _add_event( self, - data: KVCacheCreatedData | KVCacheStoredData | KVCacheRemovedData | KVCacheUpdatedData, + data: (KVCacheCreatedData | KVCacheStoredData | KVCacheRemovedData | KVCacheUpdatedData), layer_group_id: LayerGroupId = None, ) -> None: if self._max_kv_event_entries <= 0: @@ -416,7 +416,7 @@ def _flush_all_removed_events_unlocked(self) -> None: def _add_event_unlocked( self, - data: KVCacheCreatedData | KVCacheStoredData | KVCacheRemovedData | KVCacheUpdatedData, + data: (KVCacheCreatedData | KVCacheStoredData | KVCacheRemovedData | KVCacheUpdatedData), layer_group_id: LayerGroupId = None, ) -> KVCacheEvent: if not isinstance(data, KVCacheRemovedData): @@ -524,12 +524,6 @@ def _drop_hash_cache(self, block_hash: bytes) -> None: self._v1_hash_compatible_keys.discard(block_hash) self._v1_root_attrs_by_block_key.pop(block_hash, None) - @staticmethod - def _resolve_page_ref(page_ref: Any) -> Any: - if page_ref is None: - return None - return page_ref() if callable(page_ref) else page_ref - @staticmethod def _normalize_token(token: TokenIdExt) -> UniqueToken: if isinstance(token, bytes): @@ -542,13 +536,11 @@ def _stored_block_from_radix_block( cache_level: CacheLevel = GPU_LEVEL priority: Priority = PRIORITY_DEFAULT found_page = False - for life_cycle_id, page_ref in enumerate(block.storage): + for life_cycle_id in range(len(block.storage)): if life_cycle_ids is not None and life_cycle_id not in life_cycle_ids: continue - if page_ref is None: - continue - page = self._resolve_page_ref(page_ref) - if page is None: + page = block.get_page(life_cycle_id) + if page is None or page.num_tokens_in_block < len(block.tokens): continue cache_level = page.cache_level priority = page.priority @@ -568,11 +560,12 @@ def _stored_block_from_radix_block( @staticmethod def _life_cycle_ids_from_radix_block(block: Any) -> set[int]: - return { - life_cycle_id - for life_cycle_id, page_ref in enumerate(block.storage) - if page_ref is not None and KVCacheEventManager._resolve_page_ref(page_ref) is not None - } + life_cycle_ids = set[int]() + for life_cycle_id in range(len(block.storage)): + page = block.get_page(life_cycle_id) + if page is not None and page.num_tokens_in_block >= len(block.tokens): + life_cycle_ids.add(life_cycle_id) + return life_cycle_ids def _parent_hash_from_radix_block(self, block: Any) -> EventBlockHash | None: parent = block.prev diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py index 4a5bdc44394f..f8f5e3214f6e 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py @@ -49,6 +49,30 @@ def active_page_stats(kv_cache: Any) -> tuple[list[int], list[int]]: return counts, unscheduled_evictable +def committed_page_is_linked(kv_cache: Any, ordinal: int, lc_id: int) -> bool | None: + """Whether the sequence's page at ``(ordinal, lc_id)`` still points at a tree block. + + ``None`` when the slot is empty or holds an uncommitted page. Test hook: in C++ the + back-pointer is raw, so a page left pointing at a block that dies first is read after + free, and freed-but-mapped memory reads back plausibly enough that only a sanitizer + build catches the fault itself. + """ + cpp_introspection = _cpp_introspection_module() + if cpp_introspection is not None: + return cpp_introspection.committed_page_is_linked(kv_cache, ordinal, lc_id) + + from ._common import DEFAULT_BEAM_INDEX + from ._page import CommittedPage + + block_page = kv_cache._page(ordinal, DEFAULT_BEAM_INDEX, lc_id) + if block_page is None: + return None + page = block_page.page + if not isinstance(page, CommittedPage): + return None + return page.block() is not None + + def all_tree_pages_droppable(manager: Any) -> bool: """Return whether every page reachable from the radix tree is droppable.""" cpp_introspection = _cpp_introspection_module() @@ -237,8 +261,9 @@ def reuse_match_pages( """Match ``tokens`` against the radix tree and report reusable pages per block. Returns ``(num_tokens, pages)`` where ``pages[i]`` is ``None`` when block ``i`` - holds no page for lifecycle ``lc_id``, otherwise ``(slot_id, num_tokens_in_block)`` - with ``num_tokens_in_block`` set only for SSM pages (``None`` for attention pages). + holds no page for lifecycle ``lc_id``, otherwise ``(slot_id, num_tokens_in_block)``. + See ``CommittedPage.num_tokens_in_block`` for how attention and SSM life cycles + interpret the recorded token count. """ cpp_introspection = _cpp_introspection_module() if cpp_introspection is not None: @@ -254,16 +279,13 @@ def reuse_match_pages( pages.append((slot_id, None if num_tokens_in_block < 0 else num_tokens_in_block)) return num_tokens, pages - from ._utils import unwrap_rawref - match = manager._radix_tree.match(reuse_scope, list(tokens), enable_partial) py_pages: list[tuple[int, int | None] | None] = [] for block in match.blocks: - ref = block.storage[lc_id] - if ref is None: + page = block.get_page(lc_id) + if page is None: py_pages.append(None) else: - page = unwrap_rawref(ref) py_pages.append((page.slot_id, getattr(page, "num_tokens_in_block", None))) return match.num_tokens, py_pages @@ -286,13 +308,11 @@ def reuse_match_planned_drop_counts( manager, reuse_scope, list(tokens), lc_id, enable_partial ) - from ._utils import unwrap_rawref - match = manager._radix_tree.match(reuse_scope, list(tokens), enable_partial) counts: list[int | None] = [] for block in match.blocks: - ref = block.storage[lc_id] - counts.append(None if ref is None else unwrap_rawref(ref).planned_drop_count) + page = block.get_page(lc_id) + counts.append(None if page is None else page.planned_drop_count) return match.num_tokens, counts diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_page.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_page.py index 70a82c090c01..c3cee03270f5 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_page.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_page.py @@ -164,48 +164,36 @@ def __init__( ) self.set_slot(slot) - def convert_to_committed(self, block: Block, ready_event: CachedCudaEvent) -> "CommittedPage": + def convert_to_committed( + self, block: Block, ready_event: CachedCudaEvent, num_tokens_in_block: int + ) -> "CommittedPage": """ Moves the slot to a new committed page and add the new page to the block. The uncommitted page becomes invalid. + + `num_tokens_in_block` records the page's token count. See + `CommittedPage.num_tokens_in_block` for its attention and SSM interpretations. """ assert not self.scheduled_for_eviction - assert block.storage[self.life_cycle] is None + # Check before building: replace_page() below drops the superseded page, so a + # failure in between must not lose a usable snapshot. + assert block.can_replace_page(self.life_cycle, num_tokens_in_block) # If you hit this assertion failure, it's likely because you are using debugpy, which delayed GC # for _KVCache._take_uncommitted_page(). Disable breakpoints on exceptions to avoid this issue. assert self.status == PageStatus.DROPPABLE, "Release holder/lock first" self.ready_event = ready_event committed_page = CommittedPage( - self.manager, block, self.life_cycle, self.cache_level, self, self.priority - ) - assert not self.has_valid_slot and self.ready_event is CachedCudaEvent.NULL - assert committed_page.has_valid_slot - block.storage[self.life_cycle] = rawref.ref(committed_page) - return committed_page - - def convert_to_ssm_committed( - self, block: Block, ready_event: CachedCudaEvent, num_tokens_in_block: int - ) -> "SsmCommittedPage": - """ - Moves the slot to a new committed SSM page and add the new page to the block. - The uncommitted page becomes invalid. - """ - assert not self.scheduled_for_eviction - assert block.storage[self.life_cycle] is None - assert self.status == PageStatus.DROPPABLE, "Release holder/lock first" - self.ready_event = ready_event - committed_page = SsmCommittedPage( self.manager, block, self.life_cycle, self.cache_level, self, - self.priority, num_tokens_in_block, + self.priority, ) assert not self.has_valid_slot and self.ready_event is CachedCudaEvent.NULL assert committed_page.has_valid_slot - block.storage[self.life_cycle] = rawref.ref(committed_page) + block.replace_page(self.life_cycle, committed_page) return committed_page def __del__(self) -> None: @@ -241,6 +229,16 @@ class CommittedPage(Page): """ block: rawref.ref["Block"] + # Token count recorded for this page. It is usually len(block.tokens), but a snapshot + # taken at an earlier token boundary may live in a block that spans more tokens -- see + # Block.__init__ and _KVCache._snapshot_partial_block_to_tree. + # + # Attention and SSM life cycles interpret it differently: + # * for attention pages, it is the number of leading tokens with valid per-token KV, + # so the page is reusable for any prefix up to that count (compare with `>=`); + # * for an SSM page, it is the exact recurrent-state checkpoint, so reuse must be + # truncated to exactly that boundary. + num_tokens_in_block: int planned_drop_count: int __rawref__: rawref.ref["CommittedPage"] @@ -254,9 +252,12 @@ def __init__( life_cycle: LifeCycleId, cache_level: CacheLevel, slot: Slot, + num_tokens_in_block: int, priority: Priority, ): + assert 0 < num_tokens_in_block <= len(block.tokens) self.block = rawref.ref(block) + self.num_tokens_in_block = num_tokens_in_block self.planned_drop_count = 0 self.__rawref__ = rawref.NULL Page.__init__( @@ -276,7 +277,7 @@ def __del__(self) -> None: block = self.block() # block may be None when rebase happens, i.e. another block with the same key is committed, # replacing it, but the page is still used by a _KVCache. - if block is not None and block.unlink_page(self.life_cycle, self): + if block is not None and block.unlink_page(self.life_cycle, self) is not None: Block.clear_stale_blocks_after_page_unlink( block, self.life_cycle, @@ -286,25 +287,6 @@ def __del__(self) -> None: self.__rawref__.invalidate() -@dataclass(slots=True) -class SsmCommittedPage(CommittedPage): - num_tokens_in_block: int - - def __init__( - self, - storage: "StorageManager", - block: Block, - life_cycle: LifeCycleId, - cache_level: CacheLevel, - slot: Slot, - priority: Priority, - num_tokens_in_block: int, - ): - assert num_tokens_in_block > 0 - self.num_tokens_in_block = num_tokens_in_block - CommittedPage.__init__(self, storage, block, life_cycle, cache_level, slot, priority) - - @dataclass(slots=True) class _PageHolder: "Prevents pages from being dropped." @@ -329,7 +311,10 @@ def __del__(self) -> None: if not page.scheduled_for_eviction: page.manager.schedule_for_eviction(page) block = page.block() - if block is None or block.is_orphan: + # A page that no longer sits in its block's slot (orphaned block, or replaced + # by a page with a larger recorded token count) is unreachable for reuse, so + # keeping it in the eviction LRU would just pin a slot until memory pressure hits. + if block is None or block.is_orphan or not block.holds_page(page): page.manager.exclude_from_eviction(page) elif page.scheduled_for_eviction: page = cast(UncommittedPage, self.page) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py index c3d88dc0433b..576d400451b0 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py @@ -87,7 +87,8 @@ class _FakePage: - def __init__(self, cache_level=_DEFAULT_CACHE_LEVEL, priority=0): + def __init__(self, num_tokens_in_block, cache_level=_DEFAULT_CACHE_LEVEL, priority=0): + self.num_tokens_in_block = num_tokens_in_block self.cache_level = cache_level self.priority = priority @@ -118,7 +119,11 @@ def __init__(self, key, tokens, num_life_cycles=1, prev=None): self.tokens = tokens self.prev = prev or _FakeRootBlock() self.ordinal = getattr(self.prev, "ordinal", -1) + 1 - self.storage = [_FakePageRef(_FakePage()) for _ in range(num_life_cycles)] + self.storage = [_FakePageRef(_FakePage(len(tokens))) for _ in range(num_life_cycles)] + + def get_page(self, lc_idx): + page_ref = self.storage[lc_idx] + return None if page_ref is None else page_ref() with temporary_sys_path(os.path.dirname(os.path.abspath(__file__))): @@ -1133,6 +1138,21 @@ def test_v2_kv_cache_event_manager_readds_life_cycle_emits_stored_event(): assert "layer_groups" not in events[0]["data"] +def test_v2_kv_cache_event_manager_omits_partial_life_cycle_coverage(): + event_manager = KVCacheEventManager(max_kv_event_entries=8, window_size=128) + block = _FakeBlock(b"\xab\xd3", [1, 2], num_life_cycles=2) + block.storage[1] = _FakePageRef(_FakePage(num_tokens_in_block=1)) + + event_manager.add_stored_block_event_from_block(block) + events = _flush_serialized_events(event_manager) + + assert [event["layer_group_id"] for event in events] == [0] + assert [token["token_id"] for token in events[0]["data"]["blocks"][0]["tokens"]] == [1, 2] + + event_manager.add_stored_life_cycle_event_from_block(block, 1) + assert _flush_serialized_events(event_manager) == [] + + def test_v2_kv_cache_event_manager_reemits_stored_after_all_life_cycles_were_removed(): event_manager = KVCacheEventManager(max_kv_event_entries=8, window_size=128) block = _FakeBlock(b"\xab\xd0", [1, 2]) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py index e5b17c6bece3..4c93a13c58d5 100755 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py +++ b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py @@ -26,7 +26,7 @@ from importlib.util import find_spec from random import randbytes from statistics import median -from typing import TYPE_CHECKING, Iterator, NamedTuple, cast, get_type_hints +from typing import TYPE_CHECKING, Any, Iterator, NamedTuple, cast, get_type_hints if not TYPE_CHECKING and find_spec("kv_cache_manager_v2") is not None: from kv_cache_manager_v2 import ( @@ -148,8 +148,19 @@ from kernels import HostGate, enable_kernel_delay +KV_CACHE_MANAGER_V2_BACKEND = os.environ.get("TLLM_KV_CACHE_MANAGER_V2_BACKEND", "cpp").lower() + +# Gate for white-box tests that reach into the pure-Python implementation's objects +# (e.g. mutating a CommittedPage field). Prefer `_introspection`, which works on both +# backends; use this only when the behaviour under test cannot be reached through it. +requires_python_backend = unittest.skipIf( + KV_CACHE_MANAGER_V2_BACKEND == "cpp", + "white-box test over pure-Python KVCacheManagerV2 internals", +) + + def get_cached_cuda_event_type(): - backend = os.environ.get("TLLM_KV_CACHE_MANAGER_V2_BACKEND", "cpp").lower() + backend = KV_CACHE_MANAGER_V2_BACKEND if backend == "cpp": try: from bindings.internal.batch_manager.kv_cache_manager_v2 import CachedCudaEvent @@ -745,6 +756,43 @@ def plan_drop(tokens: list[TokenIdExt]) -> PlannedDropHandle: with self.assertRaisesRegex(ValueError, "already been dropped"): long_handle.drop() + @requires_python_backend + def test_planned_drop_handle_rejects_partial_coverage(self) -> None: + # plan_committed_block_drop() rejects via _prune_match, which clamps the match to + # the page's recorded token count, so the endpoint no longer matches exactly. + # Forcing that state needs a direct write to the page, hence the backend gate. + window_size = 8 + tokens_per_block = 8 + self.prepare(16 << 20, 0, 0, 2, window_size, 0, tokens_per_block=tokens_per_block) + tokens = [self.next_token() for _ in range(3 * tokens_per_block)] + + with TemporaryCudaStream([]) as stream_holder: + stream = cast(CudaStream, stream_holder.handle) + kv_cache = self.manager.create_kv_cache(None, tokens) + self.assertTrue(kv_cache.resume(stream)) + self.assertTrue(kv_cache.resize(len(tokens))) + kv_cache.commit(tokens) + kv_cache.stop_committing() + + swa_lc_id = next( + lc_id + for lc_id, lc in self.manager._life_cycles.attention_life_cycles() + if lc.window_size is not None + ) + tree_block = kv_cache._blocks[2].tree_block + assert tree_block is not None + page = tree_block.get_page(swa_lc_id) + assert page is not None + self.assertEqual(page.num_tokens_in_block, len(tree_block.tokens)) + + page.num_tokens_in_block -= 1 + try: + self.assertIsNone(kv_cache.plan_committed_block_drop()) + finally: + page.num_tokens_in_block += 1 + kv_cache.close() + stream_holder.take_finish_event().synchronize() + def test_reuse_scope_isolates_reuse(self) -> None: self.prepare(16 << 20, 0, 0, 2, None, 0, tokens_per_block=8) tokens = [TokenId(i) for i in range(64)] @@ -2497,6 +2545,53 @@ def test_commit_is_end_moves_partial_attention_and_ssm_pages(self) -> None: self.assertEqual(ssm_page[0], ssm_slot) self.assertEqual(ssm_page[1], 16) + def test_ssm_snapshot_moves_to_covering_block(self) -> None: + """A snapshot on a partial block survives the full sibling that replaces it.""" + tokens_per_block = 32 + cfg = self._make_ssm_config(tokens_per_block=tokens_per_block, enable_partial_reuse=True) + self.manager = KVCacheManager(cfg) + stream_holder = CachedCudaStream() + stream = cast(CudaStream, stream_holder.handle) + prompt = [self.next_token() for _ in range(96)] + ssm_lc_id = _introspection.ssm_life_cycle_id(self.manager) + assert ssm_lc_id is not None + + # Turn 1 ends at 48 tokens, i.e. 16 tokens into block 1. + kv1 = self.manager.create_kv_cache() + kv1.resume(stream) + kv1.capacity = 48 + kv1.history_length = 48 + kv1.commit(prompt[:48]) + kv1.close() + + # Turn 2 fills blocks 0..2. Block 1 becomes a full 32-token block that replaces the + # 16-token one, and its own SSM snapshot is taken on block 2, not block 1 -- so + # without moving the pages over, the 48-token endpoint would be lost. + kv2 = self.manager.create_kv_cache(input_tokens=prompt) + kv2.resume(stream) + kv2.capacity = len(prompt) + kv2.history_length = len(prompt) + kv2.commit(prompt[kv2.num_committed_tokens :]) + kv2.close() + + num_tokens, pages = _introspection.reuse_match_pages( + self.manager, ReuseScope(), prompt[:48], ssm_lc_id, True + ) + self.assertEqual(num_tokens, 48) + self.assertEqual(len(pages), 2) + self.assertIsNotNone(pages[-1]) + self.assertEqual(cast(tuple, pages[-1])[1], 16) + # Block 1 is the full 32-token sibling, not the original 16-token block: the + # 96-token prompt still matches end to end through it. + full_match, _ = _introspection.reuse_match_pages( + self.manager, ReuseScope(), prompt, ssm_lc_id, True + ) + self.assertEqual(full_match, len(prompt)) + + del kv1, kv2 + gc.collect() + stream_holder.synchronize() + def test_commit_min_snapshot_requires_history_alignment(self) -> None: """commit_min_snapshot requires commit() to start or end at history length.""" cfg = self._make_ssm_config(tokens_per_block=32) @@ -3780,6 +3875,281 @@ def test_reuse_across_prefill_turns_keeps_only_window_minus_one(self) -> None: self.manager.clear_reusable_blocks() +class TestPartialCoverageReuse(TestKVCacheManagerV2): + """Pages that cover fewer tokens than their block spans. + + A rewind endpoint (a partial trailing block, e.g. 16 tokens of a 32-token block) used + to be destroyed when a longer sibling was created, or refused when it was created + second. The covering block may not have a committable page for every life cycle at + that boundary. For SWA, commit_min_snapshot releases out-of-window + pages, while scratch reuse uses temporary shared storage that is not preserved as a + per-block page. Moving the partial sibling's pages into the covering block keeps the + endpoint reusable, with each page tagged by its recorded token count. + """ + + TOKENS_PER_BLOCK = 32 + WINDOW_SIZE = 16 + + def prepare_partial(self, gpu_quota: int = 64 << 20, window_size: int | None = None) -> None: + kv_buf_size = 8192 + window_size = self.WINDOW_SIZE if window_size is None else window_size + self.cfg = KVCacheManagerConfig( + tokens_per_block=self.TOKENS_PER_BLOCK, + cache_tiers=[GpuCacheTierConfig(quota=gpu_quota)], + layers=[ + AttentionLayerConfig( + layer_id=LayerId(0), + buffers=[ + BufferConfig(role=Role.KEY, size=kv_buf_size), + BufferConfig(role=Role.VALUE, size=kv_buf_size), + ], + ), + AttentionLayerConfig( + layer_id=LayerId(1), + buffers=[ + BufferConfig(role=Role.KEY, size=kv_buf_size), + BufferConfig(role=Role.VALUE, size=kv_buf_size), + ], + sliding_window_size=window_size, + ), + ], + enable_partial_reuse=True, + commit_min_snapshot=True, + ) + self.engine = FakeEngine(self.cfg) + self.manager = KVCacheManager(self.cfg) + + @property + def _full_attn_lc_id(self) -> int: + swa = set(_introspection.swa_life_cycle_ids(self.manager)) + return next( + lc_id + for lc_id in _introspection.attention_life_cycle_ids(self.manager) + if lc_id not in swa + ) + + @property + def _swa_lc_id(self) -> int: + return _introspection.swa_life_cycle_ids(self.manager)[0] + + def run_turn(self, prompt: list[TokenIdExt], refcheck: bool = False) -> int: + """Reuse what we can, prefill the rest, commit, close. Returns the reused count. + + `refcheck` runs FakeEngine, which validates the reused history KV against the + expected values for every layer group, so a clean run proves the reused pages hold + real data rather than uninitialized memory. It requires `resize()` to declare the + pre-prefill history length so that every block the engine writes is in the SWA + window; without it the manager is told the final length up front and skips + allocating out-of-window blocks -- which is exactly the situation that leaves a + partial page behind, so the two cannot be combined in the same turn. + """ + with TemporaryCudaStream([]) as s: + stream = cast(CudaStream, s.handle) + kv_cache = self.manager.create_kv_cache(input_tokens=prompt) + num_reused = kv_cache.num_committed_tokens + self.assertTrue(kv_cache.resume(stream)) + self.assertTrue(kv_cache.resize(len(prompt), num_reused if refcheck else len(prompt))) + history = list(prompt[:num_reused]) + inp = list(prompt[num_reused:]) + if refcheck: + self.engine.execute([Step(kv_cache, inp, history)], stream) + kv_cache.commit(inp) + kv_cache.close() + s.take_finish_event().synchronize() + return num_reused + + def _tail_coverage(self, prompt: list[TokenIdExt], lc_id: int) -> int: + """Recorded token count of `lc_id`'s page on the last block matching `prompt`. + + Zero when that block holds no page for the life cycle. This is the tree block + holding the tail of `prompt` (block 2 in these tests). + """ + _, pages = _introspection.reuse_match_pages(self.manager, ReuseScope(), prompt, lc_id, True) + self.assertTrue(pages) + page = pages[-1] + if page is None: + return 0 + self.assertIsNotNone(page[1]) + return cast(int, page[1]) + + def _assert_page_unlinked(self, kv_cache: Any, ordinal: int) -> None: + """Every attention page the sequence holds at `ordinal` must be off the tree. + + A page pushed out of its slot keeps no pointer to the block, so nothing + dereferences that block once it dies. + """ + for lc_id in _introspection.attention_life_cycle_ids(self.manager): + self.assertIs( + _introspection.committed_page_is_linked(kv_cache, ordinal, lc_id), + False, + f"page at ordinal {ordinal} lc {lc_id} still points at a block", + ) + + def test_replaced_page_does_not_outlive_its_block(self) -> None: + """A page pushed out of a slot must not keep a pointer to a block that dies first. + + turn1 ends inside block 2 and stays open, holding the 8-token pages it committed. + turn2 grows that block to 16 tokens and commits over them; turn3 replaces the + block with a 32-token one and destroys it. + """ + self.prepare_partial() + tpb = self.TOKENS_PER_BLOCK + prompt = [TokenId(i) for i in range(3 * tpb)] + turn1, turn2 = prompt[: 2 * tpb + 8], prompt[: 2 * tpb + 16] + + with TemporaryCudaStream([]) as s: + stream = cast(CudaStream, s.handle) + kv1 = self.manager.create_kv_cache(input_tokens=turn1) + self.assertTrue(kv1.resume(stream)) + self.assertTrue(kv1.resize(len(turn1), len(turn1))) + kv1.commit(turn1[kv1.num_committed_tokens :], is_end=True) + # kv1 stays open, holding block 2's 8-token pages. + + kv2 = self.manager.create_kv_cache(input_tokens=turn2) + self.assertEqual(kv2.num_committed_tokens, len(turn1)) + self.assertTrue(kv2.resume(stream)) + self.assertTrue(kv2.resize(len(turn2), len(turn2))) + # Block 2 becomes a 16-token block that commits over kv1's 8-token pages, + # pushing them out of the slot while kv1 still holds them. + kv2.commit(turn2[kv2.num_committed_tokens :], is_end=True) + kv2.close() + + try: + # kv1's 8-token pages left the slot but kv1 still holds them. + self._assert_page_unlinked(kv1, 2) + + # A 32-token block replaces the 16-token one, which is then destroyed. + self.run_turn(prompt) + self._assert_page_unlinked(kv1, 2) + finally: + # Must run even on failure: leaving a sequence open makes teardown abort + # the process, which would bury the assertion message. + # Dereferences CommittedPage::block for the surviving 8-token pages. + kv1.close() + s.take_finish_event().synchronize() + + def test_replaced_reused_page_does_not_outlive_its_block(self) -> None: + """Same defect as above, reached through reuse rather than through commit. + + The holder matches the 8-token endpoint and holds the tree's pages. It must not + resume: resume()'s deferred copy swaps a reused partial page for a private one and + drops the holder. + """ + self.prepare_partial() + tpb = self.TOKENS_PER_BLOCK + prompt = [TokenId(i) for i in range(3 * tpb)] + short, mid = prompt[: 2 * tpb + 8], prompt[: 2 * tpb + 16] + + self.assertEqual(self.run_turn(short), 0) + + with TemporaryCudaStream([]) as s: + holder = self.manager.create_kv_cache(input_tokens=short) + self.assertEqual(holder.num_committed_tokens, len(short)) + + # Block 2 grows to 16 tokens, adopting the 8-token pages and then committing + # over them while `holder` still holds them. + try: + self.assertEqual(self.run_turn(mid), len(short)) + self._assert_page_unlinked(holder, 2) + # A 32-token block replaces the 16-token one, which is then destroyed. + self.assertEqual(self.run_turn(prompt), len(mid)) + self._assert_page_unlinked(holder, 2) + finally: + # Must run even on failure -- see the note in the sibling test. + holder.close() + s.take_finish_event().synchronize() + + def test_rewind_endpoint_survives_longer_sibling_created_after(self) -> None: + self.prepare_partial() + base = [TokenId(i) for i in range(80)] + extended = base + [TokenId(i) for i in range(1000, 1080)] + rewind = base + [TokenId(2000)] + + self.assertEqual(self.run_turn(base), 0) + self.assertEqual(self.run_turn(extended), len(base)) + # The 16-token endpoint block is gone, but its SWA page moved into the full + # 32-token sibling and still covers the first 16 tokens. Full coverage of + # TOKENS_PER_BLOCK also proves the block now spans a whole block, since a page + # never records more tokens than its block holds. + self.assertEqual(self._tail_coverage(rewind, self._full_attn_lc_id), self.TOKENS_PER_BLOCK) + self.assertEqual( + self._tail_coverage(rewind, self._swa_lc_id), len(base) % self.TOKENS_PER_BLOCK + ) + self.assertEqual(self.manager.probe_reuse(input_tokens=rewind), len(base)) + # The partial SWA page is stale at the longer endpoint and must not constrain + # the full-attention lifecycle's reusable prefix. + self.assertEqual(self.manager.probe_reuse(input_tokens=extended), len(extended)) + + def test_rewind_endpoint_attaches_to_longer_sibling_created_before(self) -> None: + self.prepare_partial() + base = [TokenId(i) for i in range(80)] + extended = base + [TokenId(i) for i in range(1000, 1080)] + rewind = base + [TokenId(2000)] + + self.assertEqual(self.run_turn(extended), 0) + # Block 2 already spans 32 tokens, so the 80-token prompt cannot create its own + # 16-token endpoint block; its partial pages are attached to the longer sibling. + self.assertEqual(self.run_turn(base), 0) + self.assertEqual(self.manager.probe_reuse(input_tokens=rewind), len(base)) + + def test_reused_partial_coverage_kv_is_correct(self) -> None: + """The salvaged endpoint must hold real KV, not uninitialized memory.""" + self.prepare_partial() + base = [TokenId(i) for i in range(80)] + extended = base + [TokenId(i) for i in range(1000, 1080)] + rewind = base + [TokenId(3000 + i) for i in range(40)] + + self.run_turn(base, refcheck=True) + self.run_turn(extended) + # FakeEngine checks every reused history token of both layer groups, including the + # 16 tokens of block 2 that only the salvaged partial SWA page covers. + self.assertEqual(self.run_turn(rewind, refcheck=True), len(base)) + + def test_exact_boundary_ignores_stale_last_block_partial_coverage(self) -> None: + self.prepare_partial(window_size=1) + base = [TokenId(i) for i in range(80)] + boundary = base + [TokenId(i) for i in range(1000, 1016)] + + self.assertEqual(self.run_turn(base), 0) + self.assertEqual(self.run_turn(boundary), len(base)) + + self.assertEqual( + self._tail_coverage(boundary, self._full_attn_lc_id), self.TOKENS_PER_BLOCK + ) + self.assertEqual(self._tail_coverage(boundary, self._swa_lc_id), 16) + # At the 96-token boundary the input token is the entire size-1 SWA window, so no + # historical SWA block is active. The partial SWA page must not constrain the full + # attention lifecycle's reusable prefix. + self.assertEqual(self.manager.probe_reuse(input_tokens=boundary), len(boundary)) + + def test_page_coverage_only_grows(self) -> None: + self.prepare_partial() + base = [TokenId(i) for i in range(80)] + longer_partial = base + [TokenId(i) for i in range(1000, 1008)] # 88 tokens + rewind = base + [TokenId(2000)] + + self.run_turn(base) + self.assertEqual(self._tail_coverage(rewind, self._swa_lc_id), 16) + + self.run_turn(longer_partial) + # A slot keeps only the widest page. The 24-token snapshot supersedes the 16-token + # one, and it still covers the shorter rewind endpoint. + self.assertEqual(self._tail_coverage(rewind, self._swa_lc_id), 24) + self.assertEqual(self.manager.probe_reuse(input_tokens=rewind), len(base)) + self.assertEqual( + self.manager.probe_reuse(input_tokens=longer_partial + [TokenId(2000)]), + len(longer_partial), + ) + + # A later shorter snapshot cannot replace the wider page. + self.assertEqual(self.run_turn(base), len(base)) + self.assertEqual(self._tail_coverage(rewind, self._swa_lc_id), 24) + self.assertEqual( + self.manager.probe_reuse(input_tokens=longer_partial + [TokenId(2000)]), + len(longer_partial), + ) + + class TestSlotAllocatorShrink(unittest.TestCase): def test_shrink_underused_pool(self) -> None: # Regression for NVBug 6225866: shrinking a pool whose new size is