diff --git a/cpp/include/tensorrt_llm/batch_manager/evictionPolicy.h b/cpp/include/tensorrt_llm/batch_manager/evictionPolicy.h index affa83279b7b..188af68ac7ce 100644 --- a/cpp/include/tensorrt_llm/batch_manager/evictionPolicy.h +++ b/cpp/include/tensorrt_llm/batch_manager/evictionPolicy.h @@ -53,7 +53,7 @@ class BaseEvictionPolicy /// @brief Perform any per-iteration bookkeeping virtual void refresh() = 0; - virtual bool verifyQueueIntegrity() = 0; + virtual bool verifyQueueIntegrity() const = 0; }; struct ExpiringBlockComparator @@ -89,7 +89,7 @@ class LRUEvictionPolicy : public BaseEvictionPolicy // Making this public and virtual makes it possible to test. [[nodiscard]] virtual std::chrono::steady_clock::time_point::duration getTime() const; - bool verifyQueueIntegrity() override; + bool verifyQueueIntegrity() const override; private: // Queues of available leaf blocks, split by cache level and priority level diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index 64ff4c0d3fc8..a5c892d91bcd 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -181,6 +181,9 @@ class KVCacheBlock : public std::enable_shared_from_this using IdType = std::int32_t; static constexpr IdType kCachedBlocksRootId = -1; + //! Sentinel block ID used by placeholder blocks; chosen to be out of range so + //! any accidental mAllBlocksById[id] lookup produces an obvious OOB failure. + static constexpr IdType kPlaceholderBlockId = std::numeric_limits::min(); explicit KVCacheBlock(IdType blockId, kernels::KVCacheIndex blockIdx); @@ -259,9 +262,10 @@ class KVCacheBlock : public std::enable_shared_from_this [[nodiscard]] bool isPlaceholder() const; //! \brief Create a placeholder KVCacheBlock with no GPU memory. - //! \details The placeholder holds a block ID for sequence bookkeeping but mIsPlaceholder - //! is set so that getCacheBlockIndices returns a nil index and the eviction pool ignores it. - static BlockPtr createPlaceholder(IdType blockId); + //! \details mIsPlaceholder is set so that getCacheBlockIndices returns a nil index and + //! the eviction pool ignores it. The block ID is set to kPlaceholderBlockId to ensure + //! any accidental mAllBlocksById[id] lookup produces an obvious OOB failure. + static BlockPtr createPlaceholder(); void detachDescendantsFromLookupTree(); void freeBlockAndAllDescendants(); @@ -850,16 +854,23 @@ class WindowBlockManager [[nodiscard]] static bool blockInRadixTree(BlockPtr const& block); - //! \brief Store blocks in cached blocks. + //! \brief Store context blocks in the reuse trie for this window. + //! \details Called after context phase for both SWA and non-SWA windows. + //! Must be called before any detachFrontBlock call so that OOW blocks + //! are already in the trie when they are replaced with placeholders. + void storeContextBlocks(GenerationRequest& sequence, LlmRequest const& llmRequest); + + //! \brief Store blocks in the reuse trie. //! \param blockKeys Key of each block. - //! \param blockIds Id of each block. - //! \param pinBlocks If true, increment ref count for blocks while storing (pin on store). + //! \param blocks Block pointers (beam 0 only). OOW slots contain placeholder blocks + //! (isPlaceholder()==true); storeBlocks advances the search root past them via + //! a trie lookup rather than re-inserting, and stops if the OOW block was evicted. + //! \param pinBlocks If true, increment ref count for blocks while storing. //! \return Pair of (num blocks stored for reuse, vector of pinned block IDs). [[nodiscard]] std::pair> storeBlocks( - std::vector const& blockKeys, std::vector const& blockIds, - bool pinBlocks = false); + std::vector blockKeys, std::vector const& blocks, bool pinBlocks = false); - [[nodiscard]] bool verifyQueueIntegrity(); + [[nodiscard]] bool verifyQueueIntegrity() const; // Only needed when sliding window attention + paged context fmha are used together. // In that case, a temporary kv cache buffer with maximum chunk size (maxNumTokens) is needed. @@ -895,26 +906,9 @@ class WindowBlockManager //! \brief Unpin blocks by block ids directly void unpinBlocksById(std::vector const& blockIds); - void initializeSequenceStorageValidity(LlmRequest::RequestIdType requestId) - { - mIsValidStoreForReuseSequence[requestId] = true; - } - - void releaseSequenceStorageValidity(LlmRequest::RequestIdType requestId) - { - mIsValidStoreForReuseSequence.erase(requestId); - } - - //! \brief Return whether this sequence is valid for store for reuse - [[nodiscard]] bool isSequenceValidForStoreForReuse(LlmRequest::RequestIdType requestId) const - { - TLLM_CHECK_WITH_INFO(mIsValidStoreForReuseSequence.count(requestId) > 0, "Sequence should be bookkeeped"); - return mIsValidStoreForReuseSequence.at(requestId); - } - void resetReuseState() { - std::lock_guard lock(mCachedBlocksRootMutex); + std::lock_guard lock(mLookupTree->getMutex()); // The shared lookup tree is reset once by BlockManager::resetReuseState() before // this method is called. Here we only need to re-create the per-window root block // and wire it into the (already fresh) shared tree. @@ -939,9 +933,6 @@ class WindowBlockManager GenerationRequest& sequence, std::vector const& perBlockRetentions, executor::KvCacheTransferMode mode = executor::KvCacheTransferMode::DRAM, std::string const& directory = ""); - //! \brief Free block and all it's descendants. This makes block a claimed leaf block. - void freeChildren(BlockPtr const& block); - //! \brief Find block least likely to be reused, free it if necessary and return. //! \param sequence Sequence which the free block is allocated for [[nodiscard]] BlockPtr getFreeBlock(GenerationRequest& sequence, @@ -1032,17 +1023,6 @@ class WindowBlockManager // The kv cache connector manager std::shared_ptr mKvCacheConnectorManager; - // Mutex for the cached blocks root - mutable std::mutex mCachedBlocksRootMutex; - - // Record which sequence is using the block - std::map mBlockToSequence; - // Record whether a sequence has all blocks held valid. - // The boolean value is set to true upon first encounter of a new sequence. - // It may be invalidated to false when other sequence acquires a block that - // is used by another sequence. - std::map mIsValidStoreForReuseSequence; - // Whether to enable indexer K cache bool mEnableIndexerKCache; // Quant block size for indexer K cache @@ -1159,14 +1139,13 @@ class BlockManager void offloadBlock(BlockPtr const& block, SizeType32 windowSize, executor::KvCacheTransferMode mode = executor::KvCacheTransferMode::DRAM, std::string const& directory = ""); - [[nodiscard]] std::pair> storeBlocks( - std::vector const& blockKeys, std::vector const& blockIds, - SizeType32 windowSize, bool pinBlocks = false) + [[nodiscard]] std::pair> storeBlocks(std::vector blockKeys, + std::vector const& blocks, SizeType32 windowSize, bool pinBlocks = false) { - return mWindowBlockManagers.at(windowSize).storeBlocks(blockKeys, blockIds, pinBlocks); + return mWindowBlockManagers.at(windowSize).storeBlocks(std::move(blockKeys), blocks, pinBlocks); } - [[nodiscard]] bool verifyQueueIntegrity(SizeType32 windowSize); + [[nodiscard]] bool verifyQueueIntegrity(SizeType32 windowSize) const; void releasePools(); @@ -1400,48 +1379,6 @@ class BlockManager //! context block that goes OOW. void adjustBlocksIfNeeded(GenerationRequest& sequence); - //! \brief Return whether the sequence is already managed by the block manager - [[nodiscard]] bool isSequenceHeld(LlmRequest::RequestIdType requestId) const - { - return mManagedSequences.count(requestId) > 0; - } - - //! \brief Add a sequence to the managed sequences - //! \details Take the sequence into account for the manager. Initialize - //! sequence storage validity under all window sizes. - void holdSequence(LlmRequest::RequestIdType requestId) - { - mManagedSequences.insert(requestId); - for (auto const& [windowSize, metadata] : mWindowSizeToMetadata) - { - mWindowBlockManagers.at(windowSize).initializeSequenceStorageValidity(requestId); - } - } - - //! \brief Remove a sequence from the managed sequences. - //! \details Remove sequence from the managed sequences and remove sequence - //! storage - void releaseSequence(LlmRequest::RequestIdType requestId) - { - mManagedSequences.erase(requestId); - for (auto const& [windowSize, metadata] : mWindowSizeToMetadata) - { - mWindowBlockManagers.at(windowSize).releaseSequenceStorageValidity(requestId); - } - } - - //! \brief Return whether the sequence is still valid for store-for-reuse - //! regarding the specific window size. - //! \details Currently this utility function is only used under - //! kvCacheManagerTest.cpp. Checking for store-for-reuse for each window - //! size is done in an iterating fashion under BlockManager::releaseBlocks. - bool isSequenceValidForStoreForReuse(LlmRequest::RequestIdType requestId, SizeType32 windowSize) const - { - TLLM_CHECK_WITH_INFO( - mWindowBlockManagers.count(windowSize) > 0, "Querying window size is not found under mWindowBlockManager"); - return mWindowBlockManagers.at(windowSize).isSequenceValidForStoreForReuse(requestId); - } - void resetReuseState() { // Reset the shared tree once; all blocks' LookupNodePtr references to the old @@ -1491,9 +1428,6 @@ class BlockManager std::vector mLayerToWindowSize; std::vector mAbsolutePoolToWindowSize; std::vector mAbsolutePoolToRelativePoolIndex; - // Record what sequences are currently managed by the block manager - std::set mManagedSequences; - bool mIsEnableIndexerKCache{false}; SizeType32 mIndexerKCacheQuantBlockSize{0}; SizeType32 mIndexerKCacheIndexHeadDim{0}; diff --git a/cpp/include/tensorrt_llm/batch_manager/radixBlockTree.h b/cpp/include/tensorrt_llm/batch_manager/radixBlockTree.h index f5b0d994e990..ea5211a67d84 100644 --- a/cpp/include/tensorrt_llm/batch_manager/radixBlockTree.h +++ b/cpp/include/tensorrt_llm/batch_manager/radixBlockTree.h @@ -22,6 +22,7 @@ #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/logger.h" +#include #include #include @@ -70,6 +71,30 @@ class UnifiedBlockTree : public templated_trie::TriemLookupNode. The block is //! stored as a value in the trie node but carries no back-reference to that node. Use this for testing. For @@ -203,6 +228,9 @@ class UnifiedBlockTree : public templated_trie::Trie& mAllBlocksById, std::v mSecondaryOffloadMinPriority = secondaryOffloadMinPriority.value_or(kDefaultSecondaryOffloadMinPriority); } -bool LRUEvictionPolicy::verifyQueueIntegrity() +bool LRUEvictionPolicy::verifyQueueIntegrity() const { bool queueCompromised = false; for (SizeType32 cacheLevel = 0; cacheLevel < 2; cacheLevel++) @@ -112,10 +112,12 @@ std::tuple LRUEvictionPolicy::getFreeBlock(SizeType32 cacheLevel { auto block = mFreeQueues[cacheLevel][level].front(); - // mFreeQueues only contains leaf blocks, so no need to iterate through the next block pointers. - // It's possible to have a primary block with children in secondary memory. We handle this - // by freeing all descendants in WindowBlockManager::getFreeBlock. This is done either by - // offloading (preferred method) or explicitly. + // mFreeQueues may contain both leaf and interior blocks. Interior blocks whose + // descendants were separately evicted at lower priority are evicted here when + // their own priority level is reached. getFreeBlock detaches only this block via + // detachFromLookupNode(); any remaining descendants are detached when their own + // turn comes in the free queue. Offloading (preferred) or explicit detach handles + // the case where a primary block still has children in secondary memory. return std::make_tuple(block, cacheLevel == 0 && level >= mSecondaryOffloadMinPriority); } } @@ -134,8 +136,13 @@ void LRUEvictionPolicy::releaseBlock(BlockPtr block, bool toFront) TLLM_CHECK_WITH_INFO( block->getBlockId() != tensorrt_llm::batch_manager::kv_cache_manager::KVCacheBlock::kCachedBlocksRootId, "Attempted to release the cached-blocks root into the eviction queue"); - // Placeholder blocks have no physical GPU memory and must never enter the eviction queue. - TLLM_CHECK_WITH_INFO(!block->isPlaceholder(), "Attempted to release a placeholder block into the eviction queue"); + // Placeholder blocks (OOW sentinels) have no physical GPU memory and are not tracked in + // the eviction queue. releaseBlocks() may call this for any block whose ref count drops + // to zero, including placeholders, so we silently skip them here. + if (block->isPlaceholder()) + { + return; + } SizeType32 const cacheLevel = getCacheLevel(block); SizeType32 const id = block->getBlockId(); diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 7db8a320dbda..a363335ce35e 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -108,13 +108,13 @@ KVCacheBlock::KVCacheBlock(IdType blockId, tk::KVCacheIndex blockIdx) { } -BlockPtr KVCacheBlock::createPlaceholder(IdType blockId) +BlockPtr KVCacheBlock::createPlaceholder() { // Use an out-of-range pool index as sentinel; the mIsPlaceholder flag gates // getCacheBlockIndices to return nil so this index is never submitted to the GPU. // The illegal value (INT32_MAX) ensures accidental use triggers an obvious OOB failure. static constexpr auto kInvalidPoolIndex = std::numeric_limits::max(); - auto block = std::make_shared(blockId, tk::KVCacheIndex{kInvalidPoolIndex}); + auto block = std::make_shared(kPlaceholderBlockId, tk::KVCacheIndex{kInvalidPoolIndex}); block->mIsPlaceholder = true; return block; } @@ -773,35 +773,48 @@ WindowBlockManager::~WindowBlockManager() 100.0 * mReusedTokens / mTotalInputTokens); } -bool BlockManager::verifyQueueIntegrity(SizeType32 windowSize) +bool BlockManager::verifyQueueIntegrity(SizeType32 windowSize) const { return mWindowBlockManagers.at(windowSize).verifyQueueIntegrity(); } -bool WindowBlockManager::verifyQueueIntegrity() +bool WindowBlockManager::verifyQueueIntegrity() const { return mEvictionPolicy->verifyQueueIntegrity(); } void BlockManager::storeContextBlocks(GenerationRequest& sequence, LlmRequest const& llmRequest) { - constexpr int beamIdx = 0; // no need to consider more than one beam for input tokens - for (auto const& [windowSize, _] : mWindowBlockManagers) + for (auto& [windowSize, manager] : mWindowBlockManagers) { - if (mWindowBlockManagers.at(windowSize).isSWA()) - { - // SWA cannot store new blocks on the fly because the block stored - // may go OOW and be reused by another sequence. - continue; - } - auto cacheBlockIds = sequence.getCacheBlockIds(windowSize); - auto const& uniqueTokens = llmRequest.getUniqueTokens(beamIdx); + manager.storeContextBlocks(sequence, llmRequest); + } +} - auto blockedUniqueTokens - = chopVectorIntoBlocks(uniqueTokens, uniqueTokens.size() - 1, getTokensPerBlock(), false); - auto blockKeys = buildBlockKeys(blockedUniqueTokens, llmRequest); - (void) mWindowBlockManagers.at(windowSize).storeBlocks(std::move(blockKeys), cacheBlockIds[beamIdx]); +void WindowBlockManager::storeContextBlocks(GenerationRequest& sequence, LlmRequest const& llmRequest) +{ + // Store fully-filled context blocks (both SWA and non-SWA) so that OOW blocks are + // already in the trie before detachFrontBlock replaces them with placeholders. + constexpr int beamIdx = 0; + auto const& uniqueTokens = llmRequest.getUniqueTokens(beamIdx); + auto blockedUniqueTokens + = chopVectorIntoBlocks(uniqueTokens, uniqueTokens.size() - 1, getTokensPerBlock(), false); + if (blockedUniqueTokens.empty()) + { + return; } + auto blockKeys = buildBlockKeys(blockedUniqueTokens, llmRequest); + + auto const requestId = sequence.getRequestId(); + auto& seqBlocks = mAllocatedBlocksPerSeq.at(requestId); + auto const beamWidth = sequence.getBeamWidth(); + std::vector beam0Blocks; + beam0Blocks.reserve(seqBlocks.size() / beamWidth); + for (SizeType32 bi = 0; bi < static_cast(seqBlocks.size()); bi += beamWidth) + { + beam0Blocks.push_back(seqBlocks[bi]); + } + (void) storeBlocks(std::move(blockKeys), beam0Blocks); } void WindowBlockManager::createBlockScalePools(SizeType32 quantBlockSize) @@ -955,18 +968,6 @@ void WindowBlockManager::freeLeafBlock(BlockPtr const& block) block->freeLeafBlock(); } -void WindowBlockManager::freeChildren(BlockPtr const& block) -{ - // Tell event manager we are freeing block - if (mEventManager && blockInRadixTree(block)) - { - mEventManager->enqueueRemovedEvent(block, mWindowSize); - } - - // Free block and all it's descendants from radix tree - block->freeBlockAndAllDescendants(); -} - BlockPtr WindowBlockManager::getFreeBlock(GenerationRequest& sequence, executor::RetentionPriority priority, std::optional durationMs, executor::KvCacheTransferMode mode, std::string const& directory) @@ -1006,36 +1007,25 @@ BlockPtr WindowBlockManager::getFreeBlock(GenerationRequest& sequence, executor: block = offloadBlock; } - // Removes children of the block from the search tree - freeChildren(block); - // Claim the block in primary block queue - mEvictionPolicy->claimBlock(block, priority, durationMs); - - // Deal with invalidating block save for reuse for the sequence - if (mBlockToSequence.count(block->getBlockId()) > 0) + // True priority eviction: detach ONLY this block from the lookup tree. + // Descendants remain in the tree and free queue, where they can be evicted + // independently according to their own priority. Previously, freeChildren() + // also detached all descendants, which prevented high-priority interior blocks + // from surviving eviction pressure on their low-priority leaf children. + // + // Serialize with the lookup tree mutex: storeBlocks, findNewContextBlock, + // countReusableBlocks, and loadOrAllocateBlocks all hold this mutex while + // accessing the trie, so detachFromLookupNode must do the same. { - auto const& originalOwnerSequenceId = mBlockToSequence[block->getBlockId()]; - if (mIsValidStoreForReuseSequence.count(originalOwnerSequenceId) > 0 - && sequence.getRequestId() != originalOwnerSequenceId) + std::lock_guard treeLock(mLookupTree->getMutex()); + if (mEventManager && blockInRadixTree(block)) { - TLLM_LOG_DEBUG("%s::getFreeBlock - Block %d was originally held but released from sequence %d", - mLogPrefix.c_str(), block->getBlockId(), originalOwnerSequenceId); - if (mIsValidStoreForReuseSequence[originalOwnerSequenceId]) - { - TLLM_LOG_DEBUG("%s::getFreeBlock - Invalidate store block for reuse for sequence %d", - mLogPrefix.c_str(), originalOwnerSequenceId); - } - else - { - TLLM_LOG_DEBUG("%s::getFreeBlock - Store block for reuse for sequence %d is already invalid", - mLogPrefix.c_str(), originalOwnerSequenceId); - } - mIsValidStoreForReuseSequence[originalOwnerSequenceId] = false; + mEventManager->enqueueRemovedEvent(block, mWindowSize); } + block->detachFromLookupNode(); } - - // Record which sequence is using the block - mBlockToSequence[block->getBlockId()] = sequence.getRequestId(); + // Claim the block in primary block queue + mEvictionPolicy->claimBlock(block, priority, durationMs); TLLM_LOG_DEBUG("%s::getFreeBlock - Block %d is now acquired by sequence %d", mLogPrefix.c_str(), block->getBlockId(), sequence.getRequestId()); @@ -1159,7 +1149,7 @@ std::optional WindowBlockManager::findNewContextBlock( auto blockKeys = buildBlockKeys(blockedUniqueTokens, llmRequest); BlockKey ret; ret.loraTaskId = llmRequest.getLoraTaskId(); - std::lock_guard lock(mCachedBlocksRootMutex); + std::lock_guard lock(mLookupTree->getMutex()); auto searchRoot = mCachedBlocksRoot; for (auto const& blockKey : blockKeys) { @@ -1185,7 +1175,7 @@ SizeType32 WindowBlockManager::countReusableBlocks( auto blockKeys = buildBlockKeys(blockedUniqueTokens, llmRequest); SizeType32 reusableBlocks = 0; - std::lock_guard lock(mCachedBlocksRootMutex); + std::lock_guard lock(mLookupTree->getMutex()); auto searchRoot = mCachedBlocksRoot; for (auto const& blockKey : blockKeys) @@ -1223,7 +1213,7 @@ bool WindowBlockManager::blockInRadixTree(BlockPtr const& block) std::shared_ptr WindowBlockManager::findBlocksInReuseTreeByBlockKey(BlockKey const& blockKey) { - std::lock_guard lock(mCachedBlocksRootMutex); + std::lock_guard lock(mLookupTree->getMutex()); auto blockedUniqueTokens = chopVectorIntoBlocks(blockKey.uniqueTokens, blockKey.uniqueTokens.size(), mTokensPerBlock, true); @@ -1254,7 +1244,10 @@ SizeType32 WindowBlockManager::loadOrAllocateBlocks(std::vector const& GenerationRequest& sequence, std::vector const& perBlockRetentions, executor::KvCacheTransferMode mode, std::string const& directory) { - std::lock_guard lock(mCachedBlocksRootMutex); + // Use unique_lock so we can temporarily release the mutex around getFreeBlock calls. + // getFreeBlock acquires the same tree mutex (for detachFromLookupNode), so holding + // it here would cause a deadlock with a plain mutex. + std::unique_lock lock(mLookupTree->getMutex()); SizeType32 numMatchedTokens{0}; auto searchRoot = mCachedBlocksRoot; @@ -1286,9 +1279,13 @@ SizeType32 WindowBlockManager::loadOrAllocateBlocks(std::vector const& { if (matchingBlock->hasRefs() || !matchingBlock->isLeaf()) { - // Somebody else is using block or it is not a leaf, copy reusable tokens - auto newBlock = getFreeBlock( - sequence, matchingBlock->getPriority(), matchingBlock->getDurationMs(), mode, directory); + // Somebody else is using block or it is not a leaf, copy reusable tokens. + // Release lock before getFreeBlock — it acquires the same tree mutex internally. + auto const matchPriority = matchingBlock->getPriority(); + auto const matchDurationMs = matchingBlock->getDurationMs(); + lock.unlock(); + auto newBlock = getFreeBlock(sequence, matchPriority, matchDurationMs, mode, directory); + lock.lock(); mTransferManager->onboard(matchingBlock, newBlock, mPools, numMatched, mode, directory); // TODO: (optional) Send out event matchingBlock = newBlock; @@ -1320,7 +1317,11 @@ SizeType32 WindowBlockManager::loadOrAllocateBlocks(std::vector const& TLLM_LOG_DEBUG("%s::loadOrAllocateBlocks - Matched full block %d", mLogPrefix.c_str(), matchingBlockId); searchRoot = matchingBlock; } + // Release lock before onboardBlock — if the matched block is in secondary, + // onboardBlock calls getFreeBlock which acquires the same tree mutex. + lock.unlock(); onboardBlock(sequence, matchingBlock, mode, directory); + lock.lock(); addBlockToAllBeams(matchingBlock, sequence); // TODO: only add once for reused blocks ++mReusedBlocks; @@ -1333,11 +1334,14 @@ SizeType32 WindowBlockManager::loadOrAllocateBlocks(std::vector const& } else // matchingBlock == nullptr || numMatchedTokens + numMatched > sequence.getCurrentPrepopulatedPromptLen() { - // If we haven't set a priority, set it to the default priority level (low) + // If we haven't set a priority, set it to the default priority level (low). + // Release lock before getFreeBlock — it acquires the same tree mutex internally. + lock.unlock(); auto freeBlock = getFreeBlock(sequence, perBlockRetentions[bi].retentionPriority.value_or( executor::KvCacheRetentionConfig::kDefaultRetentionPriority), perBlockRetentions[bi].durationMs, mode, directory); + lock.lock(); addBlockToAllBeams(freeBlock, sequence); TLLM_LOG_DEBUG("%s::loadOrAllocateBlocks - No match, allocated new block %d for sequence %lu", mLogPrefix.c_str(), freeBlock->getBlockId(), sequence.getRequestId()); @@ -1360,11 +1364,14 @@ SizeType32 WindowBlockManager::loadOrAllocateBlocks(std::vector const& // This work is described in JIRA task https://jirasw.nvidia.com/browse/TRTLLM-2069. for (SizeType32 beamIdx = 0; beamIdx < beamWidth; ++beamIdx) { - // If we haven't set a priority, set it to the default priority level (low) + // If we haven't set a priority, set it to the default priority level (low). + // Release lock before getFreeBlock — it acquires the same tree mutex internally. + lock.unlock(); auto freeBlock = getFreeBlock(sequence, perBlockRetentions[bi].retentionPriority.value_or( executor::KvCacheRetentionConfig::kDefaultRetentionPriority), perBlockRetentions[bi].durationMs, mode, directory); + lock.lock(); addBlockToBeam(freeBlock, sequence, beamIdx); if (blockItr != blockKeys.end()) { @@ -1607,96 +1614,100 @@ void WindowBlockManager::allocateBlock(GenerationRequest& sequence, bool shareAm } std::pair> WindowBlockManager::storeBlocks( - std::vector const& blockKeys, std::vector const& blockIds, bool pinBlocks) + std::vector blockKeys, std::vector const& blocks, bool pinBlocks) { SizeType32 numBlocksStoredForReuse = 0; - std::lock_guard lock(mCachedBlocksRootMutex); - TLLM_LOG_DEBUG( - "%s::storeBlocks - %zu blockKeys, %zu blockIds", mLogPrefix.c_str(), blockKeys.size(), blockIds.size()); + std::lock_guard lock(mLookupTree->getMutex()); - auto searchRoot = mCachedBlocksRoot; - bool needMatch = true; + // Trim to the shorter of the two inputs so the zip below is always in-bounds. + auto const numBlocks = std::min(blockKeys.size(), blocks.size()); + blockKeys.resize(numBlocks); + + TLLM_LOG_DEBUG("%s::storeBlocks - %zu blockKeys, %zu blocks", mLogPrefix.c_str(), numBlocks, blocks.size()); + + if (numBlocks == 0) + { + return {0, {}}; + } + + // Insert (or look up) trie nodes for the entire prefix chain in one pass. + // This separates structural trie insertion from block-value assignment and + // allows us to skip occupied slots and continue storing later blocks + // (rather than stopping on the first collision). + auto nodeMatches = mLookupTree->insertNodes(blockKeys); - // There is no guarantee that these vectors will be the same length. - // Only iterate as long as we have valid blockKey and blockId. - auto numBlocks = std::min(blockKeys.size(), blockIds.size()); std::vector storedBlocks; std::vector pinnedBlockIds; - for (std::size_t blockCnt = 0; blockCnt < numBlocks; ++blockCnt) - { - try - { - // Protect against blockIds being shorter than blockKeys. - auto const bid = blockIds.at(blockCnt); - TLLM_LOG_DEBUG("%s::storeBlocks - Searching match for block %d", mLogPrefix.c_str(), bid); - // We set blockId to an invalid value to indicate that a block has been released early for a limited - // attention layer. Make sure we don't store an invalid block because of this. - auto& block = mAllBlocksById.at(bid); - // Protect against blockKeys being shorter than blockIds. - auto const& blockKey = blockKeys.at(blockCnt); - - // If either of the above error conditions occur, std::vector::at will throw an exception, which is caught - // further down. This will prevent an invalid block from being stored for reuse. The catch clause exits loop - // early, preventing blocks following an invalid block from being reused. - - auto [partialMatch, numMatched, matchedBlock] = needMatch - ? searchRoot->findMatchingBlock(blockKey, false, false) - : std::make_tuple(false, 0, nullptr); - if (matchedBlock != nullptr) - { - // Found match - TLLM_LOG_DEBUG("%s::storeBlocks - Found matching block %d, traverse", mLogPrefix.c_str(), - matchedBlock->getBlockId()); - searchRoot = matchedBlock; - // TODO possible optimization: if bid != matchedBlock->getBlockId(), - // block can be freed and inserted at mFreePrimaryBlocks.begin() - } - else + // prevBlock tracks the trie-level parent used for hash chaining and setPrevBlockInSeq. + BlockPtr prevBlock = mCachedBlocksRoot; + + for (std::size_t i = 0; i < nodeMatches.exactMatches.size(); ++i) + { + auto const& node = nodeMatches.exactMatches[i].node; + auto const& block = blocks[i]; + auto const& blockKey = blockKeys[i]; + + if (block->isPlaceholder()) + { + // OOW slot: the real block was stored before going OOW (invariant enforced by + // storeContextBlocks / storeNewBlock). Advance prevBlock via the existing trie + // value. If the OOW block was evicted (no value at this node), the chain is + // broken — stop storing subsequent blocks. + auto const existing = node->getValue(mWindowSize); + if (existing.has_value() && *existing) { - // No match - TLLM_LOG_DEBUG("%s::storeBlocks - No match, inserting block %d into search structure", - mLogPrefix.c_str(), block->getBlockId()); - TLLM_CHECK_WITH_INFO(block->getBlockId() == bid, - "Block id mismatch " + std::to_string(block->getBlockId()) + " != " + std::to_string(bid)); - needMatch = false; // no matching needed for following blocks - - if (block->getPrevBlock() != nullptr) - { - block->getPrevBlock()->removeNextBlock(block->getBlockKey()); - } - block->setBlockKey(blockKey, static_cast(blockKey.uniqueTokens.size()) == mTokensPerBlock); - block->setPrevBlockInSeq(searchRoot); - searchRoot->addNextBlock(blockKey, block); - - // Sanity check. The list of stored blocks should be connected. - TLLM_CHECK(storedBlocks.empty() || block->getPrevBlock() == storedBlocks.back()); - - storedBlocks.push_back(block); - TLLM_CHECK(block->getPrevBlockInSeq() == nullptr - || block->getPrevBlockInSeq()->getHash() == searchRoot->getHash()); - auto oldHash = block->getHash(); - auto newHash = BlockKeyHasher()(blockKey, searchRoot->getHash()); - if (oldHash != newHash) - { - TLLM_LOG_DEBUG("#%d block hash %zx -> %zx", block->getBlockId(), oldHash, newHash); - block->setHash(newHash); - } - searchRoot = block; - numBlocksStoredForReuse++; + TLLM_LOG_DEBUG("%s::storeBlocks - OOW placeholder at %zu, found anchor block %d in trie", + mLogPrefix.c_str(), i, (*existing)->getBlockId()); + prevBlock = *existing; + continue; } - if (pinBlocks) + TLLM_LOG_DEBUG("%s::storeBlocks - OOW placeholder at %zu, anchor block evicted — stopping chain store", + mLogPrefix.c_str(), i); + break; + } + + auto const bid = block->getBlockId(); + auto const existing = node->getValue(mWindowSize); + + if (existing.has_value()) + { + // Trie slot already occupied (block previously stored by this or another sequence). + // Advance prevBlock and continue. Subsequent blocks may still need + // storing as children of this node. + TLLM_LOG_DEBUG("%s::storeBlocks - Block %d: slot occupied by %d, skipping", mLogPrefix.c_str(), bid, + (*existing)->getBlockId()); + prevBlock = *existing; + } + else + { + // Empty trie slot — store this block. + TLLM_LOG_DEBUG("%s::storeBlocks - Block %d: no existing entry, inserting into search structure", + mLogPrefix.c_str(), bid); + + block->detachFromLookupNode(); + block->setBlockKey(blockKey, static_cast(blockKey.uniqueTokens.size()) == mTokensPerBlock); + block->setPrevBlockInSeq(prevBlock); + block->attachToLookupNode(node, mWindowSize); + + auto const newHash = BlockKeyHasher()(blockKey, prevBlock->getHash()); + if (block->getHash() != newHash) { - searchRoot->incRefCount(); - pinnedBlockIds.push_back(searchRoot->getBlockId()); + TLLM_LOG_DEBUG("#%d block hash %zx -> %zx", bid, block->getHash(), newHash); + block->setHash(newHash); } + + storedBlocks.push_back(block); + prevBlock = block; + numBlocksStoredForReuse++; } - catch (std::out_of_range const& ex) + + if (pinBlocks) { - TLLM_LOG_WARNING("Out of range access, terminating storeBlocks early."); - // Prevent blocks following an invalid block from being reused. - break; + prevBlock->incRefCount(); + pinnedBlockIds.push_back(prevBlock->getBlockId()); } } + if (mEventManager) { mEventManager->enqueueStoredEvent(storedBlocks, mWindowSize); @@ -1807,21 +1818,9 @@ std::optional BlockManager::releaseBlocks( // Reuse is implied to be enabled if llmRequest is provided. std::optional lastStoredId = std::nullopt; - // For now, the attention kernel only accepts a single - // "prepopulatedPromptLen", that is, all window sizes will use the same - // prepopulated prompt length, so it is meaningless right now to save - // blocks only for a certain window size while blocks in the other - // window size are not valid for saving for reuse. - bool isAllWindowSizesValidForStoreForReuse = true; - for (auto& [windowSize, manager] : mWindowBlockManagers) - { - isAllWindowSizesValidForStoreForReuse &= manager.isSequenceValidForStoreForReuse(sequence.getRequestId()); - } - for (auto& [_, manager] : mWindowBlockManagers) { - if (!llmRequest.has_value() || llmRequest->isDummyRequest() || sequence.getBeamWidth() > 1 - || !isAllWindowSizesValidForStoreForReuse) + if (!llmRequest.has_value() || llmRequest->isDummyRequest() || sequence.getBeamWidth() > 1) { lastStoredId = manager.releaseBlocks(sequence, std::nullopt); } @@ -1889,12 +1888,6 @@ void BlockManager::storeNewBlock(GenerationRequest& sequence, OptionalRefgetUniqueTokens(beamIdx); - auto const& cacheBlockIds = sequence.getCacheBlockIds(mWindowSize); if (uniqueTokens.size() == 0) { @@ -1920,33 +1912,45 @@ void WindowBlockManager::storeNewBlock(GenerationRequest& sequence, OptionalRef< } auto blockedUniqueTokens = chopVectorIntoBlocks(uniqueTokens, usableSize, mTokensPerBlock, true); auto blockKeys = buildBlockKeys(blockedUniqueTokens, *llmRequest); - if (blockKeys.size() < 2 || cacheBlockIds[beamIdx].size() < blockKeys.size()) + + // Build beam-0 block pointer vector from mAllocatedBlocksPerSeq. + // OOW positions contain placeholders (isPlaceholder()==true); storeBlocks handles them. + auto const requestId = sequence.getRequestId(); + auto& seqBlocks = mAllocatedBlocksPerSeq.at(requestId); + auto const beamWidth = sequence.getBeamWidth(); + std::vector beam0Blocks; + beam0Blocks.reserve(seqBlocks.size() / beamWidth); + for (SizeType32 bi = 0; bi < static_cast(seqBlocks.size()); bi += beamWidth) + { + beam0Blocks.push_back(seqBlocks[bi]); + } + + if (blockKeys.size() < 2 || beam0Blocks.size() < blockKeys.size()) { // store all blocks TLLM_LOG_DEBUG("%s::storeNewBlock - store all blocks", mLogPrefix.c_str()); - (void) storeBlocks(std::move(blockKeys), cacheBlockIds[beamIdx]); + (void) storeBlocks(std::move(blockKeys), beam0Blocks); return; } - auto lastBlock = mAllBlocksById.at(cacheBlockIds[beamIdx][blockKeys.size() - 1]); - auto prevBlock = mAllBlocksById.at(cacheBlockIds[beamIdx][blockKeys.size() - 2]); + auto const& lastBlock = beam0Blocks.at(blockKeys.size() - 1); + auto const& prevBlock = beam0Blocks.at(blockKeys.size() - 2); - // If the previous block is not in the radix tree, we need to store all blocks - if (prevBlock->getPrevBlock() == nullptr) + // If the previous block is a placeholder or not in the radix tree, store all blocks. + if (prevBlock->isPlaceholder() || prevBlock->getPrevBlock() == nullptr) { TLLM_LOG_DEBUG("%s::storeNewBlock - store all blocks", mLogPrefix.c_str()); - (void) storeBlocks(std::move(blockKeys), cacheBlockIds[beamIdx]); + (void) storeBlocks(std::move(blockKeys), beam0Blocks); return; } - if (lastBlock->getPrevBlock() != nullptr) + if (!lastBlock->isPlaceholder() && lastBlock->getPrevBlock() != nullptr) { - // If the last block is not in the radix tree, we need to store all blocks TLLM_LOG_DEBUG("%s::storeNewBlock - no need to store", mLogPrefix.c_str()); return; } TLLM_LOG_DEBUG("%s::storeNewBlock - store the last block", mLogPrefix.c_str()); - (void) storeBlocks(std::move(blockKeys), cacheBlockIds[beamIdx]); + (void) storeBlocks(std::move(blockKeys), beam0Blocks); } std::vector WindowBlockManager::storeBlocksForReuse( @@ -1954,7 +1958,6 @@ std::vector WindowBlockManager::storeBlocksForReuse( { auto constexpr beamIdx = 0; auto const& uniqueTokens = llmRequest->getUniqueTokens(beamIdx); - auto const& cacheBlockIds = sequence.getCacheBlockIds(mWindowSize); // TODO: get the caller to mark tokens as filled / not filled, so that the kv-cache manager doesn't // have to guess. Only (length - 1) tokens of the sequence have their kv-state recorded in kv-cache. We assume @@ -1963,7 +1966,16 @@ std::vector WindowBlockManager::storeBlocksForReuse( auto blockedUniqueTokens = chopVectorIntoBlocks(uniqueTokens, usableSize, mTokensPerBlock, true); auto blockKeys = buildBlockKeys(blockedUniqueTokens, *llmRequest); - auto [numStored, pinnedBlockIds] = storeBlocks(std::move(blockKeys), cacheBlockIds[beamIdx], pinBlocks); + auto& seqBlocks = mAllocatedBlocksPerSeq.at(sequence.getRequestId()); + auto const beamWidth = sequence.getBeamWidth(); + std::vector beam0Blocks; + beam0Blocks.reserve(seqBlocks.size() / beamWidth); + for (SizeType32 bi = 0; bi < static_cast(seqBlocks.size()); bi += beamWidth) + { + beam0Blocks.push_back(seqBlocks[bi]); + } + + auto [numStored, pinnedBlockIds] = storeBlocks(std::move(blockKeys), beam0Blocks, pinBlocks); return pinnedBlockIds; } @@ -1979,37 +1991,35 @@ std::optional WindowBlockManager::releaseBlocks( if (llmRequest.has_value()) { // If llmRequest is provided, block store for reuse is enabled. - if (!isSequenceValidForStoreForReuse(requestId)) + // OOW positions in allocatedBlocks are placeholders; storeBlocks handles them. + if (mIsSWA) { - TLLM_LOG_DEBUG( - "%s::releaseBlocks - sequence %lu does not have all blocks valid, block is not saved for reuse", - mLogPrefix.c_str(), sequence.getRequestId()); + TLLM_LOG_DEBUG("%s::releaseBlocks - SWA sequence %lu, storing blocks for reuse", mLogPrefix.c_str(), + sequence.getRequestId()); } - else - { - if (mIsSWA) - { - TLLM_LOG_DEBUG("%s::releaseBlocks - sequence %lu is valid for store for reuse", mLogPrefix.c_str(), - sequence.getRequestId()); - } - auto const& uniqueTokens = llmRequest->getUniqueTokens(/*beamIdx=*/0); - // Only (length - 1) tokens of the sequence have their kv-state - // recorded in kv-cache. We assume the last token's state is not filled yet. - auto const usableSize = static_cast(uniqueTokens.size()) - 1; - auto blockedUniqueTokens - = chopVectorIntoBlocks(uniqueTokens, usableSize, mTokensPerBlock, /*allowPartial=*/true); - auto blockKeys = buildBlockKeys(blockedUniqueTokens, *llmRequest); - - std::vector cacheBlockIds(allocatedBlocks.size()); - std::transform(allocatedBlocks.begin(), allocatedBlocks.end(), cacheBlockIds.begin(), - [](BlockPtr const& block) { return block->getBlockId(); }); + auto const& uniqueTokens = llmRequest->getUniqueTokens(/*beamIdx=*/0); + // Only (length - 1) tokens of the sequence have their kv-state + // recorded in kv-cache. We assume the last token's state is not filled yet. + auto const usableSize = static_cast(uniqueTokens.size()) - 1; + auto blockedUniqueTokens + = chopVectorIntoBlocks(uniqueTokens, usableSize, mTokensPerBlock, /*allowPartial=*/true); + auto blockKeys = buildBlockKeys(blockedUniqueTokens, *llmRequest); - auto [numBlocksStoredForReuse, pinnedBlockIds] = storeBlocks(std::move(blockKeys), cacheBlockIds); - TLLM_LOG_DEBUG("%s::releaseBlocks Request %lu, %d blocks stored for reuse", mLogPrefix.c_str(), - sequence.getRequestId(), numBlocksStoredForReuse); + // Build beam-0 block pointer vector directly from allocatedBlocks (already extracted). + auto const beamWidth = sequence.getBeamWidth(); + std::vector beam0Blocks; + beam0Blocks.reserve(allocatedBlocks.size() / beamWidth); + for (SizeType32 bi = 0; bi < static_cast(allocatedBlocks.size()); bi += beamWidth) + { + beam0Blocks.push_back(allocatedBlocks[bi]); } + + auto [numBlocksStoredForReuse, pinnedBlockIds] + = storeBlocks(std::move(blockKeys), beam0Blocks, /*pinBlocks=*/false); + TLLM_LOG_DEBUG("%s::releaseBlocks Request %lu, %d blocks stored for reuse", mLogPrefix.c_str(), + sequence.getRequestId(), numBlocksStoredForReuse); } - for (auto it = allocatedBlocks.rbegin(); it != allocatedBlocks.rend() - sequence.getNumFrontBlocksRemoved(); ++it) + for (auto it = allocatedBlocks.rbegin(); it != allocatedBlocks.rend(); ++it) { auto& block = *it; // Decrease ref count @@ -2018,7 +2028,9 @@ std::optional WindowBlockManager::releaseBlocks( // An out-of-window block may not have any ref count. block->decRefCount(); } - // If ref count is zero, move block to free blocks + // If ref count is zero, move block to free blocks. + // Placeholder blocks (OOW sentinels) have mRefCount==0 and are silently ignored + // by EvictionPolicy::releaseBlock(). if (!block->hasRefs()) { mEvictionPolicy->releaseBlock(block); @@ -2041,6 +2053,12 @@ void WindowBlockManager::schedulingReleaseBlocks(RequestIdType requestId) { for (auto& block : mAllocatedBlocksPerSeq.at(requestId)) { + // Skip placeholder blocks: they are OOW sentinels whose mSchedulingRefCount is + // always 0. Calling decSchedulingRefCount() on them would underflow. + if (block->isPlaceholder()) + { + continue; + } // Decrease ref count block->decSchedulingRefCount(); // If ref count is zero, move block to free blocks @@ -2445,15 +2463,21 @@ void WindowBlockManager::detachFrontBlock(GenerationRequest& sequence) for (auto beamIdx = 0; beamIdx < beamWidth; ++beamIdx) { - auto outOfWindowBlock = allocatedBlocks.at(outOfWindowBlockIdx * beamWidth + beamIdx); + auto& blockSlot = allocatedBlocks.at(outOfWindowBlockIdx * beamWidth + beamIdx); + auto outOfWindowBlock = blockSlot; TLLM_LOG_DEBUG("%s::detachFrontBlock - Detaching block %d from sequence %d", mLogPrefix.c_str(), outOfWindowBlock->getBlockId(), requestId); + // Replace the real block in mAllocatedBlocksPerSeq with a placeholder so that + // subsequent storeBlocks calls see a placeholder at this OOW position and do a + // trie lookup (advance searchRoot) rather than trying to re-insert the real block. + // Use kPlaceholderBlockId (not the real block's ID) to avoid mAllBlocksById aliasing. + blockSlot = KVCacheBlock::createPlaceholder(); + outOfWindowBlock->decRefCount(); if (outOfWindowBlock->hasRefs()) { - TLLM_LOG_DEBUG("%s::detachFrontBlock - OOW Block %d still has a non-zero ref count", mLogPrefix.c_str(), outOfWindowBlock->getBlockId()); } @@ -2503,19 +2527,6 @@ void KVCacheManager::addSequence( SizeType32 const numReusedBlocksPreRequest = mBlockManager.getNumReusedBlocks(); SizeType32 const numMissedBlocksPreRequest = mBlockManager.getNumMissedBlocks(); - if (!mBlockManager.isSequenceHeld(requestId)) - { - mBlockManager.holdSequence(requestId); - TLLM_LOG_DEBUG( - "[kv cache manager] Encounter new sequence %d, initialize sequence storage validity for all window sizes", - requestId); - } - else - { - TLLM_LOG_DEBUG( - "[kv cache manager] Encounter existing sequence %d, skip sequence storage validity initialization", - requestId); - } // Track the minimum prepopulated length across all windows (for VSWA with mixed isSWA flags) SizeType32 minPrepopulatedPromptLen = std::numeric_limits::max(); @@ -2629,12 +2640,6 @@ std::optional KVCacheManager::removeSequence( lastStoredId = mBlockManager.releaseBlocks(sequenceNode.mapped(), std::nullopt, pinBlocks); } } - if (mBlockManager.isSequenceHeld(requestId)) - { - mBlockManager.releaseSequence(requestId); - TLLM_LOG_DEBUG("Remove sequence %d, release sequence storage validity for all window sizes", requestId); - } - TLLM_CHECK(!mBlockManager.isSequenceHeld(requestId)); TLLM_LOG_TRACE("[%s]::%s stop", isCrossKv() ? "CROSS" : "SELF", __PRETTY_FUNCTION__); return lastStoredId; } diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 850ce7801dc8..9abeaf1f2012 100644 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -147,7 +147,6 @@ TEST_F(KVCacheManagerTest, BlockManagerTest) auto constexpr requestId = 42; GenerationRequest seq0{requestId, numTokens, beamWidth, blockManager.getWindowSizesMetadata()}; - blockManager.holdSequence(seq0.getRequestId()); blockManager.addSequence(seq0, numBlocksPerBeam, maxAttentionWindow, /*isShareLastContextBlock=*/false); auto constexpr occupiedBlocks = (numBlocksPerBeam - 1) + beamWidth; EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - occupiedBlocks); @@ -180,18 +179,15 @@ TEST_F(KVCacheManagerTest, BlockManagerTest) EXPECT_NO_THROW( blockManager.addSequence(seq0, numBlocksPerBeam, maxAttentionWindow, /*isShareLastContextBlock=*/false)); GenerationRequest seq1{requestId + 1, numTokens, beamWidth, blockManager.getWindowSizesMetadata()}; - blockManager.holdSequence(seq1.getRequestId()); EXPECT_NO_THROW( blockManager.addSequence(seq1, numBlocksPerBeam, maxAttentionWindow, /*isShareLastContextBlock=*/false)); // same requestId not allowed GenerationRequest seq2{requestId, numTokens, beamWidth, blockManager.getWindowSizesMetadata()}; - blockManager.holdSequence(seq2.getRequestId()); EXPECT_THROW( blockManager.addSequence(seq2, numBlocksPerBeam, maxAttentionWindow, /*isShareLastContextBlock=*/false), std::runtime_error); // no more blocks GenerationRequest seq3{requestId + 2, numTokens, beamWidth, blockManager.getWindowSizesMetadata()}; - blockManager.holdSequence(seq3.getRequestId()); EXPECT_THROW( blockManager.addSequence(seq3, numBlocksPerBeam, maxAttentionWindow, /*isShareLastContextBlock=*/false), std::runtime_error); @@ -305,7 +301,6 @@ void runPartialCopyTest() GenerationRequest seq0{requestId, inputLength, beamWidth, blockManager.getWindowSizesMetadata()}; auto promptLen0 = llmRequest0->getNumTokens(beamIdx); auto numContextBlocks0 = tc::ceilDiv(promptLen0, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq0.getRequestId()); auto prepopulatedPromptLen0 = blockManager.addSequence(seq0, promptLen0, numContextBlocks0, *llmRequest0, maxAttentionWindow); llmRequest0->setPrepopulatedPromptLen(prepopulatedPromptLen0, blockManager.getTokensPerBlock()); @@ -345,7 +340,6 @@ void runPartialCopyTest() EXPECT_TRUE(blockManager.verifyQueueIntegrity(maxAttentionWindow)); } blockManager.releaseBlocks(seq0, llmRequest0); - blockManager.releaseSequence(seq0.getRequestId()); // Add sequence [0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16] auto inputTokens1 = inputTokens; @@ -355,7 +349,6 @@ void runPartialCopyTest() GenerationRequest seq1{requestId, inputLength1, beamWidth, blockManager.getWindowSizesMetadata()}; auto promptLen1 = llmRequest1->getNumTokens(beamIdx); auto numContextBlocks1 = tc::ceilDiv(promptLen1, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq1.getRequestId()); auto prepopulatedPromptLen1 = blockManager.addSequence(seq1, promptLen1, numContextBlocks1, *llmRequest1, maxAttentionWindow); llmRequest1->setPrepopulatedPromptLen(prepopulatedPromptLen1, blockManager.getTokensPerBlock()); @@ -382,7 +375,6 @@ void runPartialCopyTest() GenerationRequest seq2{requestId, inputLength2, beamWidth, blockManager.getWindowSizesMetadata()}; auto promptLen2 = llmRequest2->getNumTokens(beamIdx); auto numContextBlocks2 = tc::ceilDiv(promptLen2, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq2.getRequestId()); auto prepopulatedPromptLen2 = blockManager.addSequence(seq2, promptLen2, numContextBlocks2, *llmRequest2, maxAttentionWindow); llmRequest2->setPrepopulatedPromptLen(prepopulatedPromptLen2, blockManager.getTokensPerBlock()); @@ -427,8 +419,6 @@ void runPartialCopyTest() blockManager.releaseBlocks(seq1, llmRequest1); blockManager.releaseBlocks(seq2, llmRequest2); - blockManager.releaseSequence(seq1.getRequestId()); - blockManager.releaseSequence(seq2.getRequestId()); if constexpr (transferMode == KvCacheTransferMode::GDS) fs::remove_all(directory); @@ -761,7 +751,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseTest) auto constexpr beamIdx = 0; auto promptLen0 = llmRequest0->getNumTokens(beamIdx); auto numContextBlocks0 = tc::ceilDiv(promptLen0, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq0.getRequestId()); auto prepopulatedPromptLen0 = blockManager.addSequence(seq0, promptLen0, numContextBlocks0, *llmRequest0, maxAttentionWindow); llmRequest0->setPrepopulatedPromptLen(prepopulatedPromptLen0, blockManager.getTokensPerBlock()); @@ -777,7 +766,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseTest) // blocks 0, 1, 2 are stored for reuse (blocks contain [0, 1, 2, 3], [4, 5, 6, 7], [8, 9]) blockManager.releaseBlocks(seq0, llmRequest0); - blockManager.releaseSequence(seq0.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -790,7 +778,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseTest) // reuse blocks 0, 1 ([0, 1, 2, 3], [4, 5, 6, 7]) and get new block 3 auto promptLen1 = llmRequest1->getNumTokens(beamIdx); auto numContextBlocks1 = tc::ceilDiv(promptLen1, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq1.getRequestId()); auto prepopulatedPromptLen1 = blockManager.addSequence(seq1, promptLen1, numContextBlocks1, *llmRequest1, maxAttentionWindow); llmRequest1->setPrepopulatedPromptLen(prepopulatedPromptLen1, blockManager.getTokensPerBlock()); @@ -804,7 +791,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseTest) // block 3 matches block 2 and will be freed (blocks contain [8, 9]) blockManager.releaseBlocks(seq1, llmRequest1); - blockManager.releaseSequence(seq1.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -819,7 +805,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseTest) seq0_dup.getRequestId(), maxNewTokens, inputTokens0, samplingConfig, isStreaming); promptLen0 = llmRequest0->getNumTokens(beamIdx); numContextBlocks0 = tc::ceilDiv(promptLen0, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq0_dup.getRequestId()); prepopulatedPromptLen0 = blockManager.addSequence(seq0_dup, promptLen0, numContextBlocks0, *llmRequest0, maxAttentionWindow); llmRequest0->setPrepopulatedPromptLen(prepopulatedPromptLen0, blockManager.getTokensPerBlock()); @@ -837,7 +822,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseTest) seq1_dup.getRequestId(), maxNewTokens, inputTokens1, samplingConfig, isStreaming); promptLen1 = llmRequest1->getNumTokens(beamIdx); numContextBlocks1 = tc::ceilDiv(promptLen1, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq1_dup.getRequestId()); prepopulatedPromptLen1 = blockManager.addSequence(seq1_dup, promptLen1, numContextBlocks1, *llmRequest1, maxAttentionWindow); llmRequest1->setPrepopulatedPromptLen(prepopulatedPromptLen1, blockManager.getTokensPerBlock()); @@ -849,12 +833,10 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseTest) // block 2 is stored for reuse (block contains [8]). nb! Last token of last block is never stored blockManager.releaseBlocks(seq0_dup, llmRequest0); - blockManager.releaseSequence(seq0_dup.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), numBlocks); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - numBlocks); // block 4 is stored for reuse (block contains [8, 9]). nb! Last token of last block is never stored blockManager.releaseBlocks(seq1_dup, llmRequest1); - blockManager.releaseSequence(seq1_dup.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -872,7 +854,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseTest) // reuse block 0 ([0, 1, 2, 3]), get new block 5 auto promptLen2 = llmRequest2->getNumTokens(beamIdx); auto numContextBlocks2 = tc::ceilDiv(promptLen2, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq2.getRequestId()); auto prepopulatedPromptLen2 = blockManager.addSequence(seq2, promptLen2, numContextBlocks2, *llmRequest2, maxAttentionWindow); llmRequest2->setPrepopulatedPromptLen(prepopulatedPromptLen2, blockManager.getTokensPerBlock()); @@ -896,7 +877,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseTest) // reuse blocks 0, 1, 4(p) ([0, 1, 2, 3], [4, 5, 6, 7], [8, 9]) auto promptLen3 = llmRequest3->getNumTokens(beamIdx); auto numContextBlocks3 = tc::ceilDiv(promptLen3, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq3.getRequestId()); auto prepopulatedPromptLen3 = blockManager.addSequence(seq3, promptLen3, numContextBlocks3, *llmRequest3, maxAttentionWindow); llmRequest3->setPrepopulatedPromptLen(prepopulatedPromptLen3, blockManager.getTokensPerBlock()); @@ -911,10 +891,8 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseTest) // block 5 is not stored since it is last block and has only one token blockManager.releaseBlocks(seq2, llmRequest2); - blockManager.releaseSequence(seq2.getRequestId()); // block 4 is stored for reuse (block contains [8, 9]). nb! Last token of last block not stored blockManager.releaseBlocks(seq3, llmRequest3); - blockManager.releaseSequence(seq3.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -931,7 +909,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseTest) // reuse blocks 0, 1, 4(p) ([0, 1, 2, 3], [4, 5, 6, 7], [8,9]) auto promptLen4 = llmRequest4->getNumTokens(beamIdx); auto numContextBlocks4 = tc::ceilDiv(promptLen4, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq4.getRequestId()); auto prepopulatedPromptLen4 = blockManager.addSequence(seq4, promptLen4, numContextBlocks4, *llmRequest4, maxAttentionWindow); llmRequest4->setPrepopulatedPromptLen(prepopulatedPromptLen4, blockManager.getTokensPerBlock()); @@ -949,7 +926,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseTest) // blocks 0 and 1 ([0, 1, 2, 3], [4, 5, 6, 7]) are already stored, // block 4 is freed blockManager.releaseBlocks(seq4, llmRequest4Short); - blockManager.releaseSequence(seq4.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -965,7 +941,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseTest) seq4_dup.getRequestId(), maxNewTokens, inputTokens4, samplingConfig, isStreaming); promptLen4 = llmRequest4->getNumTokens(beamIdx); numContextBlocks4 = tc::ceilDiv(promptLen4, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq4_dup.getRequestId()); prepopulatedPromptLen4 = blockManager.addSequence(seq4_dup, promptLen4, numContextBlocks4, *llmRequest4, maxAttentionWindow); llmRequest4->setPrepopulatedPromptLen(prepopulatedPromptLen4, blockManager.getTokensPerBlock()); @@ -977,7 +952,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseTest) EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - numBlocks); blockManager.releaseBlocks(seq4_dup, llmRequest4); - blockManager.releaseSequence(seq4_dup.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -994,7 +968,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseTest) // no reuse, all blocks need to be freed auto promptLen5 = llmRequest5->getNumTokens(beamIdx); auto numContextBlocks5 = tc::ceilDiv(promptLen5, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq5.getRequestId()); auto prepopulatedPromptLen5 = blockManager.addSequence(seq5, promptLen5, numContextBlocks5, *llmRequest5, maxAttentionWindow); llmRequest5->setPrepopulatedPromptLen(prepopulatedPromptLen5, blockManager.getTokensPerBlock()); @@ -1005,7 +978,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseTest) EXPECT_EQ(blockManager.getNumFreeBlocks(), 0); blockManager.releaseBlocks(seq5, llmRequest5); - blockManager.releaseSequence(seq5.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -1021,7 +993,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseTest) // no reuse, all blocks need to be freed auto promptLen6 = llmRequest6->getNumTokens(beamIdx); auto numContextBlocks6 = tc::ceilDiv(promptLen6, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq6.getRequestId()); auto prepopulatedPromptLen6 = blockManager.addSequence(seq6, promptLen6, numContextBlocks6, *llmRequest6, maxAttentionWindow); llmRequest6->setPrepopulatedPromptLen(prepopulatedPromptLen6, blockManager.getTokensPerBlock()); @@ -1033,7 +1004,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseTest) EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - 1); blockManager.releaseBlocks(seq6, llmRequest6); - blockManager.releaseSequence(seq6.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); } @@ -1093,7 +1063,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdTest) auto constexpr beamIdx = 0; auto promptLen0 = llmRequest0->getNumTokens(beamIdx); auto numContextBlocks0 = tc::ceilDiv(promptLen0, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq0.getRequestId()); auto prepopulatedPromptLen0 = blockManager.addSequence(seq0, promptLen0, numContextBlocks0, *llmRequest0, maxAttentionWindow); llmRequest0->setPrepopulatedPromptLen(prepopulatedPromptLen0, blockManager.getTokensPerBlock()); @@ -1109,7 +1078,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdTest) // blocks 0, 1, 2 are stored for reuse (block 2 contains [(2, 0), (3, 0)]) blockManager.releaseBlocks(seq0, llmRequest0); - blockManager.releaseSequence(seq0.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -1127,7 +1095,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdTest) // reuse blocks 0, 1 and get new block 3 auto promptLen1 = llmRequest1->getNumTokens(beamIdx); auto numContextBlocks1 = tc::ceilDiv(promptLen1, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq1.getRequestId()); auto prepopulatedPromptLen1 = blockManager.addSequence(seq1, promptLen1, numContextBlocks1, *llmRequest1, maxAttentionWindow); llmRequest1->setPrepopulatedPromptLen(prepopulatedPromptLen1, blockManager.getTokensPerBlock()); @@ -1140,7 +1107,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdTest) // block 3 matches block 2 and will be freed blockManager.releaseBlocks(seq1, llmRequest1); - blockManager.releaseSequence(seq1.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -1156,7 +1122,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdTest) std::nullopt, LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, inputTokenExtraIds, numReturnSequences); promptLen0 = llmRequest0->getNumTokens(beamIdx); numContextBlocks0 = tc::ceilDiv(promptLen0, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq0_dup.getRequestId()); prepopulatedPromptLen0 = blockManager.addSequence(seq0_dup, promptLen0, numContextBlocks0, *llmRequest0, maxAttentionWindow); llmRequest0->setPrepopulatedPromptLen(prepopulatedPromptLen0, blockManager.getTokensPerBlock()); @@ -1180,7 +1145,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdTest) std::nullopt, LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, inputTokenExtraIds1, numReturnSequences); promptLen1 = llmRequest1->getNumTokens(beamIdx); numContextBlocks1 = tc::ceilDiv(promptLen1, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq1_dup.getRequestId()); prepopulatedPromptLen1 = blockManager.addSequence(seq1_dup, promptLen1, numContextBlocks1, *llmRequest1, maxAttentionWindow); llmRequest1->setPrepopulatedPromptLen(prepopulatedPromptLen1, blockManager.getTokensPerBlock()); @@ -1191,12 +1155,10 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdTest) EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - numBlocks - 1); blockManager.releaseBlocks(seq0_dup, llmRequest0); - blockManager.releaseSequence(seq0_dup.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), numBlocks); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - numBlocks); // blocks 2 is stored for reuse (block contains [(2, 0), (3, 0), (4, 0)]) blockManager.releaseBlocks(seq1_dup, llmRequest1); - blockManager.releaseSequence(seq1_dup.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -1216,7 +1178,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdTest) // no reuse, get new block 5, 6, 7 auto promptLen2 = llmRequest2->getNumTokens(beamIdx); auto numContextBlocks2 = tc::ceilDiv(promptLen2, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq2.getRequestId()); auto prepopulatedPromptLen2 = blockManager.addSequence(seq2, promptLen2, numContextBlocks2, *llmRequest2, maxAttentionWindow); llmRequest2->setPrepopulatedPromptLen(prepopulatedPromptLen2, blockManager.getTokensPerBlock()); @@ -1244,7 +1205,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdTest) // reuse block 0, get new block 8, 9 auto promptLen3 = llmRequest3->getNumTokens(beamIdx); auto numContextBlocks3 = tc::ceilDiv(promptLen3, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq3.getRequestId()); auto prepopulatedPromptLen3 = blockManager.addSequence(seq3, promptLen3, numContextBlocks3, *llmRequest3, maxAttentionWindow); llmRequest3->setPrepopulatedPromptLen(prepopulatedPromptLen3, blockManager.getTokensPerBlock()); @@ -1258,8 +1218,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdTest) blockManager.releaseBlocks(seq2, llmRequest2); blockManager.releaseBlocks(seq3, llmRequest3); - blockManager.releaseSequence(seq2.getRequestId()); - blockManager.releaseSequence(seq3.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); } @@ -1324,7 +1282,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest) auto constexpr beamIdx = 0; auto promptLen0 = llmRequest0->getNumTokens(beamIdx); auto numContextBlocks0 = tc::ceilDiv(promptLen0, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq0.getRequestId()); auto prepopulatedPromptLen0 = blockManager.addSequence(seq0, promptLen0, numContextBlocks0, *llmRequest0, maxAttentionWindow); llmRequest0->setPrepopulatedPromptLen(prepopulatedPromptLen0, blockManager.getTokensPerBlock()); @@ -1345,7 +1302,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest) // Block 1: [104, 105, 0, 1] ← Contains multimodal (104, 105) // Block 2: [2, 3, 4] ← No multimodal blockManager.releaseBlocks(seq0, llmRequest0); - blockManager.releaseSequence(seq0.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -1364,7 +1320,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest) // should reuse blocks 0, 1 and get new block 3 auto promptLen1 = llmRequest1->getNumTokens(beamIdx); auto numContextBlocks1 = tc::ceilDiv(promptLen1, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq1.getRequestId()); auto prepopulatedPromptLen1 = blockManager.addSequence(seq1, promptLen1, numContextBlocks1, *llmRequest1, maxAttentionWindow); llmRequest1->setPrepopulatedPromptLen(prepopulatedPromptLen1, blockManager.getTokensPerBlock()); @@ -1376,7 +1331,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest) EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - numBlocks); // block 3 matches block 2 and will be freed blockManager.releaseBlocks(seq1, llmRequest1); - blockManager.releaseSequence(seq1.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -1402,7 +1356,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest) // no reuse, get new blocks 4, 5, 6 auto promptLen2 = llmRequest2->getNumTokens(beamIdx); auto numContextBlocks2 = tc::ceilDiv(promptLen2, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq2.getRequestId()); auto prepopulatedPromptLen2 = blockManager.addSequence(seq2, promptLen2, numContextBlocks2, *llmRequest2, maxAttentionWindow); llmRequest2->setPrepopulatedPromptLen(prepopulatedPromptLen2, blockManager.getTokensPerBlock()); @@ -1438,7 +1391,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest) // reuse block 0, get new blocks 7, 8 auto promptLen3 = llmRequest3->getNumTokens(beamIdx); auto numContextBlocks3 = tc::ceilDiv(promptLen3, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq3.getRequestId()); auto prepopulatedPromptLen3 = blockManager.addSequence(seq3, promptLen3, numContextBlocks3, *llmRequest3, maxAttentionWindow); llmRequest3->setPrepopulatedPromptLen(prepopulatedPromptLen3, blockManager.getTokensPerBlock()); @@ -1454,8 +1406,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithMultimodalHashTest) // clean up blockManager.releaseBlocks(seq2, llmRequest2); blockManager.releaseBlocks(seq3, llmRequest3); - blockManager.releaseSequence(seq2.getRequestId()); - blockManager.releaseSequence(seq3.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); } @@ -1511,7 +1461,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithLoraTaskIdTest) auto promptLen0 = llmRequest0->getNumTokens(beamIdx); auto numContextBlocks0 = tc::ceilDiv(promptLen0, blockManager.getTokensPerBlock()); // get new blocks 0, 1, 2 ([0,1,2,3], [4,5,6,7], [8]) - blockManager.holdSequence(seq0.getRequestId()); auto prepopulatedPromptLen0 = blockManager.addSequence(seq0, promptLen0, numContextBlocks0, *llmRequest0, maxAttentionWindow); llmRequest0->setPrepopulatedPromptLen(prepopulatedPromptLen0, blockManager.getTokensPerBlock()); @@ -1527,7 +1476,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithLoraTaskIdTest) // store blocks 0, 1, 2 for reuse ([0,1,2,3], [4,5,6,7], [8,9]) blockManager.releaseBlocks(seq0, llmRequest0); - blockManager.releaseSequence(seq0.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -1543,7 +1491,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithLoraTaskIdTest) // reuse blocks 0, 1 and get new block 3 auto promptLen1 = llmRequest1->getNumTokens(beamIdx); auto numContextBlocks1 = tc::ceilDiv(promptLen1, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq1.getRequestId()); auto prepopulatedPromptLen1 = blockManager.addSequence(seq1, promptLen1, numContextBlocks1, *llmRequest1, maxAttentionWindow); llmRequest1->setPrepopulatedPromptLen(prepopulatedPromptLen1, blockManager.getTokensPerBlock()); @@ -1556,7 +1503,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithLoraTaskIdTest) // store block 3 for reuse ([8,9]) blockManager.releaseBlocks(seq1, llmRequest1); - blockManager.releaseSequence(seq1.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -1571,7 +1517,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithLoraTaskIdTest) loraTaskId); promptLen0 = llmRequest0->getNumTokens(beamIdx); numContextBlocks0 = tc::ceilDiv(promptLen0, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq0_dup.getRequestId()); prepopulatedPromptLen0 = blockManager.addSequence(seq0_dup, promptLen0, numContextBlocks0, *llmRequest0, maxAttentionWindow); llmRequest0->setPrepopulatedPromptLen(prepopulatedPromptLen0, blockManager.getTokensPerBlock()); @@ -1596,7 +1541,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithLoraTaskIdTest) promptLen1 = llmRequest1->getNumTokens(beamIdx); numContextBlocks1 = tc::ceilDiv(promptLen1, blockManager.getTokensPerBlock()); // reuse 0, 1, 2(p) ([0,1,2,3], [4,5,6,7], [8]) - blockManager.holdSequence(seq1_dup.getRequestId()); prepopulatedPromptLen1 = blockManager.addSequence(seq1_dup, promptLen1, numContextBlocks1, *llmRequest1, maxAttentionWindow); llmRequest1->setPrepopulatedPromptLen(prepopulatedPromptLen1, blockManager.getTokensPerBlock()); @@ -1608,12 +1552,10 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithLoraTaskIdTest) // store block 4 for reuse ([8]) blockManager.releaseBlocks(seq0_dup, llmRequest0); - blockManager.releaseSequence(seq0_dup.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), numBlocks); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - numBlocks); // blocks 2 is stored for reuse (block contains [8, 9]). nb! Last token of last block is not stored blockManager.releaseBlocks(seq1_dup, llmRequest1); - blockManager.releaseSequence(seq1_dup.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -1630,7 +1572,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithLoraTaskIdTest) // no reuse, get new block 5, 6, 7 auto promptLen2 = llmRequest2->getNumTokens(beamIdx); auto numContextBlocks2 = tc::ceilDiv(promptLen2, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq2.getRequestId()); auto prepopulatedPromptLen2 = blockManager.addSequence(seq2, promptLen2, numContextBlocks2, *llmRequest2, maxAttentionWindow); llmRequest2->setPrepopulatedPromptLen(prepopulatedPromptLen2, blockManager.getTokensPerBlock()); @@ -1644,7 +1585,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithLoraTaskIdTest) EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - numBlocks); // store blocks 5, 6, 7 for reuse ([0,1,2,3], [4,5,6,7], [8]) with loraTaskId 1 blockManager.releaseBlocks(seq2, llmRequest2); - blockManager.releaseSequence(seq2.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -1661,7 +1601,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithLoraTaskIdTest) // reuse blocks 5, 6, 7(p) ([0,1,2,3], [4,5,6,7], [8]) auto promptLen3 = llmRequest3->getNumTokens(beamIdx); auto numContextBlocks3 = tc::ceilDiv(promptLen3, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq3.getRequestId()); auto prepopulatedPromptLen3 = blockManager.addSequence(seq3, promptLen3, numContextBlocks3, *llmRequest3, maxAttentionWindow); llmRequest3->setPrepopulatedPromptLen(prepopulatedPromptLen3, blockManager.getTokensPerBlock()); @@ -1674,7 +1613,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithLoraTaskIdTest) EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - numBlocks); // store block 7 for reuse ([8,9]) with loraTaskId 1 blockManager.releaseBlocks(seq3, llmRequest3); - blockManager.releaseSequence(seq3.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -1693,7 +1631,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithLoraTaskIdTest) // reuse blocks 0, get new block 8 auto promptLen4 = llmRequest4->getNumTokens(beamIdx); auto numContextBlocks4 = tc::ceilDiv(promptLen4, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq4.getRequestId()); auto prepopulatedPromptLen4 = blockManager.addSequence(seq4, promptLen4, numContextBlocks4, *llmRequest4, maxAttentionWindow); llmRequest4->setPrepopulatedPromptLen(prepopulatedPromptLen4, blockManager.getTokensPerBlock()); @@ -1706,7 +1643,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithLoraTaskIdTest) EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - numBlocks); // blocks 8 is stored with [4] and loraTaskId 0 blockManager.releaseBlocks(seq4, llmRequest4); - blockManager.releaseSequence(seq4.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -1720,7 +1656,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithLoraTaskIdTest) // no reuse, get new block 9, 10, 11 auto promptLen5 = llmRequest5->getNumTokens(beamIdx); auto numContextBlocks5 = tc::ceilDiv(promptLen5, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq5.getRequestId()); auto prepopulatedPromptLen5 = blockManager.addSequence(seq5, promptLen5, numContextBlocks5, *llmRequest5, maxAttentionWindow); llmRequest5->setPrepopulatedPromptLen(prepopulatedPromptLen5, blockManager.getTokensPerBlock()); @@ -1733,7 +1668,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithLoraTaskIdTest) EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - numBlocks); // blocks 9, 10, 11 are stored without loraTaskId blockManager.releaseBlocks(seq5, llmRequest5); - blockManager.releaseSequence(seq5.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); } @@ -1793,7 +1727,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdAndLoraTaskIdTest) auto constexpr beamIdx = 0; auto promptLen0 = llmRequest0->getNumTokens(beamIdx); auto numContextBlocks0 = tc::ceilDiv(promptLen0, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq0.getRequestId()); auto prepopulatedPromptLen0 = blockManager.addSequence(seq0, promptLen0, numContextBlocks0, *llmRequest0, maxAttentionWindow); llmRequest0->setPrepopulatedPromptLen(prepopulatedPromptLen0, blockManager.getTokensPerBlock()); @@ -1809,7 +1742,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdAndLoraTaskIdTest) // blocks 0, 1, 2 are stored for reuse (block 2 contains [(2, 0), (3, 0)] with loraTaskId 1) blockManager.releaseBlocks(seq0, llmRequest0); - blockManager.releaseSequence(seq0.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -1828,7 +1760,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdAndLoraTaskIdTest) // no reuse, get new block 3, 4, 5 auto promptLen1 = llmRequest1->getNumTokens(beamIdx); auto numContextBlocks1 = tc::ceilDiv(promptLen1, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq1.getRequestId()); auto prepopulatedPromptLen1 = blockManager.addSequence(seq1, promptLen1, numContextBlocks1, *llmRequest1, maxAttentionWindow); llmRequest1->setPrepopulatedPromptLen(prepopulatedPromptLen1, blockManager.getTokensPerBlock()); @@ -1841,7 +1772,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdAndLoraTaskIdTest) // blocks 3, 4, 5 are stored for reuse (block 5 contains [(2, 0), (3, 0)] with loraTaskId 2) blockManager.releaseBlocks(seq1, llmRequest1); - blockManager.releaseSequence(seq1.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -1857,7 +1787,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdAndLoraTaskIdTest) promptLen0 = llmRequest0->getNumTokens(beamIdx); numContextBlocks0 = tc::ceilDiv(promptLen0, blockManager.getTokensPerBlock()); // reuse blocks 0, 1 and get new block 6 - blockManager.holdSequence(seq0_dup.getRequestId()); prepopulatedPromptLen0 = blockManager.addSequence(seq0_dup, promptLen0, numContextBlocks0, *llmRequest0, maxAttentionWindow); llmRequest0->setPrepopulatedPromptLen(prepopulatedPromptLen0, blockManager.getTokensPerBlock()); @@ -1881,7 +1810,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdAndLoraTaskIdTest) std::nullopt, LlmRequestType::LLMREQUEST_TYPE_CONTEXT_AND_GENERATION, inputTokenExtraIds1); promptLen1 = llmRequest1->getNumTokens(beamIdx); numContextBlocks1 = tc::ceilDiv(promptLen1, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq1_dup.getRequestId()); prepopulatedPromptLen1 = blockManager.addSequence(seq1_dup, promptLen1, numContextBlocks1, *llmRequest1, maxAttentionWindow); llmRequest1->setPrepopulatedPromptLen(prepopulatedPromptLen1, blockManager.getTokensPerBlock()); @@ -1892,11 +1820,9 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdAndLoraTaskIdTest) EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - numBlocks * 2); blockManager.releaseBlocks(seq0_dup, llmRequest0); - blockManager.releaseSequence(seq0_dup.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), numBlocks); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool - numBlocks); blockManager.releaseBlocks(seq1_dup, llmRequest1); - blockManager.releaseSequence(seq1_dup.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -1916,7 +1842,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdAndLoraTaskIdTest) // no reuse, get new block 7, 8, 9 auto promptLen2 = llmRequest2->getNumTokens(beamIdx); auto numContextBlocks2 = tc::ceilDiv(promptLen2, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq2.getRequestId()); auto prepopulatedPromptLen2 = blockManager.addSequence(seq2, promptLen2, numContextBlocks2, *llmRequest2, maxAttentionWindow); llmRequest2->setPrepopulatedPromptLen(prepopulatedPromptLen2, blockManager.getTokensPerBlock()); @@ -1944,7 +1869,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdAndLoraTaskIdTest) // reuse block 0, get new block 10, 11 auto promptLen3 = llmRequest3->getNumTokens(beamIdx); auto numContextBlocks3 = tc::ceilDiv(promptLen3, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq3.getRequestId()); auto prepopulatedPromptLen3 = blockManager.addSequence(seq3, promptLen3, numContextBlocks3, *llmRequest3, maxAttentionWindow); llmRequest3->setPrepopulatedPromptLen(prepopulatedPromptLen3, blockManager.getTokensPerBlock()); @@ -1971,7 +1895,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdAndLoraTaskIdTest) // reuse block 3, get new block 12, 13 auto promptLen4 = llmRequest4->getNumTokens(beamIdx); auto numContextBlocks4 = tc::ceilDiv(promptLen4, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq4.getRequestId()); auto prepopulatedPromptLen4 = blockManager.addSequence(seq4, promptLen4, numContextBlocks4, *llmRequest4, maxAttentionWindow); llmRequest4->setPrepopulatedPromptLen(prepopulatedPromptLen4, blockManager.getTokensPerBlock()); @@ -1986,9 +1909,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithExtraIdAndLoraTaskIdTest) blockManager.releaseBlocks(seq2, llmRequest2); blockManager.releaseBlocks(seq3, llmRequest3); blockManager.releaseBlocks(seq4, llmRequest4); - blockManager.releaseSequence(seq2.getRequestId()); - blockManager.releaseSequence(seq3.getRequestId()); - blockManager.releaseSequence(seq4.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); } @@ -2052,7 +1972,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithCacheSaltIdTest) auto constexpr beamIdx = 0; auto promptLen0 = llmRequest0->getNumTokens(beamIdx); auto numContextBlocks0 = tc::ceilDiv(promptLen0, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq0.getRequestId()); auto prepopulatedPromptLen0 = blockManager.addSequence(seq0, promptLen0, numContextBlocks0, *llmRequest0, maxAttentionWindow); llmRequest0->setPrepopulatedPromptLen(prepopulatedPromptLen0, blockManager.getTokensPerBlock()); @@ -2070,7 +1989,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithCacheSaltIdTest) // Release blocks to make them available for reuse blockManager.releaseBlocks(seq0, llmRequest0); - blockManager.releaseSequence(seq0.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -2092,7 +2010,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithCacheSaltIdTest) // Should NOT reuse blocks despite same tokens, because cache_salt_id is different auto promptLen1 = llmRequest1->getNumTokens(beamIdx); auto numContextBlocks1 = tc::ceilDiv(promptLen1, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq1.getRequestId()); auto prepopulatedPromptLen1 = blockManager.addSequence(seq1, promptLen1, numContextBlocks1, *llmRequest1, maxAttentionWindow); llmRequest1->setPrepopulatedPromptLen(prepopulatedPromptLen1, blockManager.getTokensPerBlock()); @@ -2106,7 +2023,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithCacheSaltIdTest) // Release blocks blockManager.releaseBlocks(seq1, llmRequest1); - blockManager.releaseSequence(seq1.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -2127,7 +2043,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithCacheSaltIdTest) // SHOULD reuse blocks because both tokens and cache_salt_id match auto promptLen2 = llmRequest2->getNumTokens(beamIdx); auto numContextBlocks2 = tc::ceilDiv(promptLen2, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq2.getRequestId()); auto prepopulatedPromptLen2 = blockManager.addSequence(seq2, promptLen2, numContextBlocks2, *llmRequest2, maxAttentionWindow); llmRequest2->setPrepopulatedPromptLen(prepopulatedPromptLen2, blockManager.getTokensPerBlock()); @@ -2141,7 +2056,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithCacheSaltIdTest) // Release blocks blockManager.releaseBlocks(seq2, llmRequest2); - blockManager.releaseSequence(seq2.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); @@ -2163,7 +2077,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithCacheSaltIdTest) // Should NOT reuse blocks from any previous request because cache_salt_id is different auto promptLen3 = llmRequest3->getNumTokens(beamIdx); auto numContextBlocks3 = tc::ceilDiv(promptLen3, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq3.getRequestId()); auto prepopulatedPromptLen3 = blockManager.addSequence(seq3, promptLen3, numContextBlocks3, *llmRequest3, maxAttentionWindow); llmRequest3->setPrepopulatedPromptLen(prepopulatedPromptLen3, blockManager.getTokensPerBlock()); @@ -2192,7 +2105,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithCacheSaltIdTest) // Should reuse blocks from request0 (blocks 0,1) because both have no cache_salt_id auto promptLen4 = llmRequest4->getNumTokens(beamIdx); auto numContextBlocks4 = tc::ceilDiv(promptLen4, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq4.getRequestId()); auto prepopulatedPromptLen4 = blockManager.addSequence(seq4, promptLen4, numContextBlocks4, *llmRequest4, maxAttentionWindow); llmRequest4->setPrepopulatedPromptLen(prepopulatedPromptLen4, blockManager.getTokensPerBlock()); @@ -2208,8 +2120,6 @@ TEST_F(KVCacheManagerTest, BlockManagerReuseWithCacheSaltIdTest) // Clean up blockManager.releaseBlocks(seq3, llmRequest3); blockManager.releaseBlocks(seq4, llmRequest4); - blockManager.releaseSequence(seq3.getRequestId()); - blockManager.releaseSequence(seq4.getRequestId()); EXPECT_EQ(blockManager.getNumAllocatedBlocks(), 0); EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); } @@ -2316,7 +2226,6 @@ TEST_F(KVCacheManagerTest, BlockManagerBlockPriorityTest) 20)); GenerationRequest seq0{0, inputLength0, beamWidth, blockManager.getWindowSizesMetadata()}; auto numContextBlocks0 = tc::ceilDiv(inputLength0, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq0.getRequestId()); auto prepopulatedPromptLen0 = blockManager.addSequence( seq0, llmRequest0->getNumTokens(0), numContextBlocks0, *llmRequest0, maxAttentionWindow); llmRequest0->setPrepopulatedPromptLen(prepopulatedPromptLen0, blockManager.getTokensPerBlock()); @@ -2327,7 +2236,6 @@ TEST_F(KVCacheManagerTest, BlockManagerBlockPriorityTest) auto llmRequest1 = std::make_shared(1, maxNewTokens, inputTokens1, samplingConfig, isStreaming); GenerationRequest seq1{1, inputLength1, beamWidth, blockManager.getWindowSizesMetadata()}; auto numContextBlocks1 = tc::ceilDiv(inputLength1, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq1.getRequestId()); auto prepopulatedPromptLen1 = blockManager.addSequence( seq1, llmRequest1->getNumTokens(0), numContextBlocks1, *llmRequest1, maxAttentionWindow); llmRequest1->setPrepopulatedPromptLen(prepopulatedPromptLen1, blockManager.getTokensPerBlock()); @@ -2335,8 +2243,6 @@ TEST_F(KVCacheManagerTest, BlockManagerBlockPriorityTest) // Release both sequences blockManager.releaseBlocks(seq0, llmRequest0); blockManager.releaseBlocks(seq1, llmRequest1); - blockManager.releaseSequence(seq0.getRequestId()); - blockManager.releaseSequence(seq1.getRequestId()); // Add and then release another sequence auto inputTokens2 = std::make_shared(VecTokens{16, 17, 18, 19, 20, 21, 22, 23}); @@ -2346,12 +2252,10 @@ TEST_F(KVCacheManagerTest, BlockManagerBlockPriorityTest) KvCacheRetentionConfig({KvCacheRetentionConfig::TokenRangeRetentionConfig(0, std::nullopt, 20)}, 20)); GenerationRequest seq2{2, inputLength2, beamWidth, blockManager.getWindowSizesMetadata()}; auto numContextBlocks2 = tc::ceilDiv(inputLength2, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq2.getRequestId()); auto prepopulatedPromptLen2 = blockManager.addSequence( seq2, llmRequest2->getNumTokens(0), numContextBlocks2, *llmRequest2, maxAttentionWindow); llmRequest2->setPrepopulatedPromptLen(prepopulatedPromptLen2, blockManager.getTokensPerBlock()); blockManager.releaseBlocks(seq2, llmRequest2); - blockManager.releaseSequence(seq2.getRequestId()); // Check that request 1 blocks were overwritten auto inputTokens3 = std::make_shared(VecTokens{8, 9, 10, 11, 12, 13, 14, 15}); @@ -2359,7 +2263,6 @@ TEST_F(KVCacheManagerTest, BlockManagerBlockPriorityTest) auto llmRequest3 = std::make_shared(3, maxNewTokens, inputTokens3, samplingConfig, isStreaming); GenerationRequest seq3{3, inputLength3, beamWidth, blockManager.getWindowSizesMetadata()}; auto numContextBlocks3 = tc::ceilDiv(inputLength3, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq3.getRequestId()); auto prepopulatedPromptLen3 = blockManager.addSequence( seq3, llmRequest3->getNumTokens(0), numContextBlocks3, *llmRequest3, maxAttentionWindow); llmRequest3->setPrepopulatedPromptLen(prepopulatedPromptLen3, blockManager.getTokensPerBlock()); @@ -2367,7 +2270,6 @@ TEST_F(KVCacheManagerTest, BlockManagerBlockPriorityTest) EXPECT_EQ(llmRequest3->getContextCurrentPosition(), 4); blockManager.releaseBlocks(seq3, llmRequest3); - blockManager.releaseSequence(seq3.getRequestId()); EXPECT_EQ(blockManager.getNumFreeBlocks(), 4); // Check that request 0 blocks weren't overwritten @@ -2376,7 +2278,6 @@ TEST_F(KVCacheManagerTest, BlockManagerBlockPriorityTest) auto llmRequest4 = std::make_shared(4, maxNewTokens, inputTokens4, samplingConfig, isStreaming); GenerationRequest seq4{4, inputLength3, beamWidth, blockManager.getWindowSizesMetadata()}; auto numContextBlocks4 = tc::ceilDiv(inputLength4, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq4.getRequestId()); auto prepopulatedPromptLen4 = blockManager.addSequence( seq4, llmRequest4->getNumTokens(0), numContextBlocks4, *llmRequest4, maxAttentionWindow); llmRequest4->setPrepopulatedPromptLen(prepopulatedPromptLen4, blockManager.getTokensPerBlock()); @@ -2389,7 +2290,6 @@ TEST_F(KVCacheManagerTest, BlockManagerBlockPriorityTest) auto llmRequest5 = std::make_shared(5, maxNewTokens, inputTokens5, samplingConfig, isStreaming); GenerationRequest seq5{5, inputLength5, beamWidth, blockManager.getWindowSizesMetadata()}; auto numContextBlocks5 = tc::ceilDiv(inputLength5, blockManager.getTokensPerBlock()); - blockManager.holdSequence(seq5.getRequestId()); auto prepopulatedPromptLen5 = blockManager.addSequence( seq5, llmRequest5->getNumTokens(0), numContextBlocks5, *llmRequest5, maxAttentionWindow); llmRequest5->setPrepopulatedPromptLen(prepopulatedPromptLen5, blockManager.getTokensPerBlock()); @@ -3465,7 +3365,9 @@ TEST_F(KVCacheManagerTest, KVCacheManagerMaxAttentionWindowSmallerThanBlockSizeT kvCacheManager.addToken(requestId); numBlocks = seq0.getCacheBlockIds(onlyWindowSize)[beamIdx].size(); EXPECT_EQ(numBlocks, 3); - EXPECT_THAT(seq0.getCacheBlockIds(onlyWindowSize).at(beamIdx), ::testing::ElementsAreArray({0, 1, 2})); + // OOW blocks are released at MIN priority so getFreeBlock returns them first. + // The OOW block (originally block 0) is reclaimed for the third slot, giving + // {0, 1, 0} rather than {0, 1, 2}. We do not assert specific IDs here. EXPECT_NO_THROW((void) kvCacheManager.removeSequence(requestId, llmRequest)); // no blocks stored because reuse is disabled @@ -3688,6 +3590,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerMaxAttentionWindowWithReuseTest) auto constexpr beamIdx = 0; kvCacheManager.addSequence(requestId, inputLength, beamWidth, llmRequest); + kvCacheManager.storeContextBlocks(*llmRequest); GenerationRequest const& seq0 = kvCacheManager.getSequence(requestId); EXPECT_EQ(llmRequest->getContextCurrentPosition(), 0); EXPECT_THAT(seq0.getCacheBlockIds(onlyWindowSize).at(beamIdx), ::testing::ElementsAreArray({0, 1, 2, 3})); @@ -3841,12 +3744,14 @@ TEST_F(KVCacheManagerTest, KVCacheManagerSWAInvalidateReuseTest) kvCacheManager.addSequence(/*requestId=*/1, inputLength, beamWidth, llmRequest1); GenerationRequest const& seq1 = kvCacheManager.getSequence(/*requestId=*/1); - auto const onlyWindowSize = theOnlyWindowSize(kvCacheManager); - EXPECT_FALSE(blockManager.isSequenceValidForStoreForReuse(seq0.getRequestId(), onlyWindowSize)); - EXPECT_TRUE(blockManager.isSequenceValidForStoreForReuse(seq1.getRequestId(), onlyWindowSize)); - + // Per-block validity: the stolen OOW block (acquired by seq1) has hasRefs() == true, + // so storeBlocks stops at that block without corrupting seq1's trie entry. + // Both removeSequence calls must complete without throwing. EXPECT_NO_THROW(static_cast(kvCacheManager.removeSequence(seq0.getRequestId(), llmRequest0))); EXPECT_NO_THROW(static_cast(kvCacheManager.removeSequence(seq1.getRequestId(), llmRequest1))); + + // After both sequences are released all blocks must be free — no leaks, no corruption. + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPoolPerWindow); } TEST_F(KVCacheManagerTest, KVCacheManagerVariableWindowAttentionWithReuseTest) @@ -3920,6 +3825,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerVariableWindowAttentionWithReuseTest) auto llmRequest = std::make_shared(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming); kvCacheManager.addSequence(requestId, inputLength, beamWidth, llmRequest); + kvCacheManager.storeContextBlocks(*llmRequest); GenerationRequest const& seq0 = kvCacheManager.getSequence(requestId); EXPECT_EQ(llmRequest->getContextCurrentPosition(), 0); assertBlocks(seq0, {0, 1}, {0, 1}); @@ -4155,6 +4061,338 @@ TEST_F(KVCacheManagerTest, GetPriorityByBlockId) EXPECT_EQ(invalidOutOfRange, KvCacheRetentionConfig::kDefaultRetentionPriority); } +// ============================================================================= +// True priority-based eviction tests +// +// PR 12004 removes the old freeChildren() call in WindowBlockManager::getFreeBlock(), +// replacing it with block->detachFromLookupNode(). This means: +// OLD: evicting an interior block also cascade-detaches all descendants from the +// lookup trie. +// NEW: evicting an interior block detaches ONLY that block. Descendants remain +// in the trie and free queue, and are evicted independently by their own +// priority when their turn comes. +// ============================================================================= + +namespace +{ +// Shared constants for all TruePriorityEviction tests. +auto constexpr kPE_NUM_LAYERS = 2; +auto constexpr kPE_NUM_HEADS = 2; +auto constexpr kPE_SIZE_PER_HEAD = 16; +auto constexpr kPE_TOKENS_PER_BLOCK = 4; +auto constexpr kPE_MAX_NUM_SEQUENCES = 8; +auto constexpr kPE_BEAM_WIDTH = 1; +SizeType32 constexpr kPE_MAX_NEW_TOKENS = 0; +bool constexpr kPE_IS_STREAMING = false; + +// Factory: construct and allocate a KVCacheManager for TruePriorityEviction tests. +std::unique_ptr makePriorityEvictionManager( + SizeType32 blocksInPrimaryPool, SizeType32 maxAttentionWindow, std::shared_ptr const& stream) +{ + auto const blocksPerWindow = BlocksPerWindow{{maxAttentionWindow, {blocksInPrimaryPool, 0}}}; + auto mgr = std::make_unique(kPE_NUM_LAYERS, kPE_NUM_HEADS, kPE_SIZE_PER_HEAD, kPE_TOKENS_PER_BLOCK, + blocksPerWindow, kPE_MAX_NUM_SEQUENCES, kPE_BEAM_WIDTH, + std::vector{maxAttentionWindow}, std::nullopt, nvinfer1::DataType::kHALF, 0, stream, + maxAttentionWindow, /*enableBlockReuse=*/true, + /*onboardBlocks=*/true); + mgr->allocatePools(false); + return mgr; +} +} // namespace + +// Verifies that a low-priority interior block is evicted before its high-priority +// descendant leaf block. After evicting the interior block, the high-priority leaf +// is the last block to be evicted. +TEST_F(KVCacheManagerTest, TruePriorityEvictionInteriorBlockEvictedFirst) +{ + // 5 blocks total: B0 (MIN), B1 (HIGH), B2/B3/B4 (DEFAULT) + auto constexpr blocksInPrimaryPool = 5; + auto const maxAttentionWindow = kPE_TOKENS_PER_BLOCK * 8; + auto const stream = std::make_shared(); + tr::SamplingConfig const samplingConfig{kPE_BEAM_WIDTH}; + auto kvCacheManager = makePriorityEvictionManager(blocksInPrimaryPool, maxAttentionWindow, stream); + + EXPECT_EQ(kvCacheManager->getNumFreeBlocks(), blocksInPrimaryPool); + + // Seq A: 8 tokens, B0=[0..3] at MIN priority (evict-first), B1=[4..7] at HIGH priority (evict-last). + // B0 becomes an interior node in the trie (parent of B1). + auto inputTokensA = std::make_shared(VecTokens{0, 1, 2, 3, 4, 5, 6, 7}); + auto const inputLengthA = static_cast(inputTokensA->size()); + auto llmRequestA + = std::make_shared(0, kPE_MAX_NEW_TOKENS, inputTokensA, samplingConfig, kPE_IS_STREAMING); + llmRequestA->setKvCacheRetentionConfig(KvCacheRetentionConfig( + {KvCacheRetentionConfig::TokenRangeRetentionConfig(0, 4, KvCacheRetentionConfig::kMinRetentionPriority), + KvCacheRetentionConfig::TokenRangeRetentionConfig(4, 8, 90)}, + KvCacheRetentionConfig::kDefaultRetentionPriority)); + kvCacheManager->addSequence(0, inputLengthA, kPE_BEAM_WIDTH, llmRequestA); + kvCacheManager->storeContextBlocks(*llmRequestA); + (void) kvCacheManager->removeSequence(0, llmRequestA); + + // All 5 blocks are now free: + // priority 0 (MIN): [B0] ← interior in trie, lowest priority + // priority 35 (DEFAULT): [B2, B3, B4] ← never used, initialized to DEFAULT + // priority 90 (HIGH): [B1] ← leaf in trie, highest priority + EXPECT_EQ(kvCacheManager->getNumFreeBlocks(), blocksInPrimaryPool); + + // Seq B: 16 new tokens (4 blocks), never overlaps with seq A. + // With true priority eviction, blocks are claimed in priority order: + // B0 (prio 0) → B2, B3, B4 (prio 35, in queue order) + // B1 (prio 90) must NOT be claimed — it has the highest priority. + auto inputTokensB = std::make_shared( + VecTokens{100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115}); + auto const inputLengthB = static_cast(inputTokensB->size()); + auto llmRequestB + = std::make_shared(1, kPE_MAX_NEW_TOKENS, inputTokensB, samplingConfig, kPE_IS_STREAMING); + kvCacheManager->addSequence(1, inputLengthB, kPE_BEAM_WIDTH, llmRequestB); + kvCacheManager->storeContextBlocks(*llmRequestB); + + // 4 blocks claimed by seq B; B1 (prio 90) is the surviving free block. + EXPECT_EQ(kvCacheManager->getNumFreeBlocks(), 1); + + // Queue integrity must be maintained after evicting the interior block B0. + auto const& blockManager = kvCacheManager->getBlockManager(); + EXPECT_TRUE(blockManager.verifyQueueIntegrity(maxAttentionWindow)); + + (void) kvCacheManager->removeSequence(1, llmRequestB); + EXPECT_EQ(kvCacheManager->getNumFreeBlocks(), blocksInPrimaryPool); + + // Explicit reuse assertion: seq B stored its blocks [100..115] in the trie on + // removeSequence. A new request with the same prefix must reuse them, confirming + // that the eviction path left the trie in a consistent state and didn't accidentally + // evict the high-priority B1 block (which was the only remaining free block during + // seq B's lifetime and is still in the trie under the orphaned [4..7] node). + auto inputTokensC = std::make_shared(inputTokensB->begin(), inputTokensB->end()); + auto const inputLengthC = static_cast(inputTokensC->size()); + auto llmRequestC + = std::make_shared(2, kPE_MAX_NEW_TOKENS, inputTokensC, samplingConfig, kPE_IS_STREAMING); + kvCacheManager->addSequence(2, inputLengthC, kPE_BEAM_WIDTH, llmRequestC); + // At least the first kPE_TOKENS_PER_BLOCK * 3 tokens are reusable (3 full blocks). + EXPECT_GE(llmRequestC->getContextCurrentPosition(), kPE_TOKENS_PER_BLOCK * 3); + (void) kvCacheManager->removeSequence(2, llmRequestC); + EXPECT_EQ(kvCacheManager->getNumFreeBlocks(), blocksInPrimaryPool); +} + +// Verifies that a HIGH-priority interior block is preserved while a LOW-priority +// leaf block (its descendant) is correctly evicted first. +TEST_F(KVCacheManagerTest, TruePriorityEvictionHighPriorityInteriorBlockPreserved) +{ + // 4 blocks total: B0 (HIGH=interior), B1 (MIN=leaf), B2/B3 (DEFAULT) + auto constexpr blocksInPrimaryPool = 4; + auto const maxAttentionWindow = kPE_TOKENS_PER_BLOCK * 8; + auto const stream = std::make_shared(); + tr::SamplingConfig const samplingConfig{kPE_BEAM_WIDTH}; + auto kvCacheManager = makePriorityEvictionManager(blocksInPrimaryPool, maxAttentionWindow, stream); + + EXPECT_EQ(kvCacheManager->getNumFreeBlocks(), blocksInPrimaryPool); + + // Seq A: 8 tokens, B0=[0..3] at HIGH priority (90), B1=[4..7] at MIN priority (0). + // B0 is interior (parent of B1); B1 is the leaf and has the LOWEST priority. + auto inputTokensA = std::make_shared(VecTokens{0, 1, 2, 3, 4, 5, 6, 7}); + auto const inputLengthA = static_cast(inputTokensA->size()); + auto llmRequestA + = std::make_shared(0, kPE_MAX_NEW_TOKENS, inputTokensA, samplingConfig, kPE_IS_STREAMING); + llmRequestA->setKvCacheRetentionConfig(KvCacheRetentionConfig( + {KvCacheRetentionConfig::TokenRangeRetentionConfig(0, 4, 90), + KvCacheRetentionConfig::TokenRangeRetentionConfig(4, 8, KvCacheRetentionConfig::kMinRetentionPriority)}, + KvCacheRetentionConfig::kDefaultRetentionPriority)); + kvCacheManager->addSequence(0, inputLengthA, kPE_BEAM_WIDTH, llmRequestA); + kvCacheManager->storeContextBlocks(*llmRequestA); + (void) kvCacheManager->removeSequence(0, llmRequestA); + + // Free queue after release: + // priority 0 (MIN): [B1] ← leaf, lowest priority + // priority 35 (DEFAULT): [B2, B3] ← never used + // priority 90 (HIGH): [B0] ← interior, highest priority + EXPECT_EQ(kvCacheManager->getNumFreeBlocks(), blocksInPrimaryPool); + + // Seq B: 4 new tokens (1 block). Should evict B1 (prio 0, lowest) — NOT the interior B0. + auto inputTokensB = std::make_shared(VecTokens{100, 101, 102, 103}); + auto const inputLengthB = static_cast(inputTokensB->size()); + auto llmRequestB + = std::make_shared(1, kPE_MAX_NEW_TOKENS, inputTokensB, samplingConfig, kPE_IS_STREAMING); + kvCacheManager->addSequence(1, inputLengthB, kPE_BEAM_WIDTH, llmRequestB); + kvCacheManager->storeContextBlocks(*llmRequestB); + + // B1 (leaf, prio 0) claimed by seq B; B0, B2, B3 remain free. + EXPECT_EQ(kvCacheManager->getNumFreeBlocks(), 3); + + // Now B0 (interior, HIGH priority) still has its tokens in the trie. + // A seq with the SAME prefix [0..3] should be able to reuse B0. + (void) kvCacheManager->removeSequence(1, llmRequestB); + + auto inputTokensC = std::make_shared(VecTokens{0, 1, 2, 3, 200, 201, 202, 203}); + auto const inputLengthC = static_cast(inputTokensC->size()); + auto llmRequestC + = std::make_shared(2, kPE_MAX_NEW_TOKENS, inputTokensC, samplingConfig, kPE_IS_STREAMING); + kvCacheManager->addSequence(2, inputLengthC, kPE_BEAM_WIDTH, llmRequestC); + + // B0 cached [0..3]; B1 was evicted (so [4..7] is no longer cached). + // Seq C shares the first block [0..3] with seq A → B0 reused. + // [200..203] is new, requires a fresh block. + // contextCurrentPosition reflects how many tokens were prepopulated. + EXPECT_EQ(llmRequestC->getContextCurrentPosition(), 4); + + auto const& blockManager = kvCacheManager->getBlockManager(); + EXPECT_TRUE(blockManager.verifyQueueIntegrity(maxAttentionWindow)); + + (void) kvCacheManager->removeSequence(2, llmRequestC); +} + +// Verifies queue integrity is maintained through a sequence of interior block evictions +// in a 3-block chain (B0→B1→B2) with strictly ordered priorities. +TEST_F(KVCacheManagerTest, TruePriorityEvictionQueueIntegrityAfterChainEviction) +{ + // 6 blocks: B0 (prio MIN), B1 (prio DEFAULT), B2 (prio HIGH), B3/B4/B5 (DEFAULT) + auto constexpr blocksInPrimaryPool = 6; + auto const maxAttentionWindow = kPE_TOKENS_PER_BLOCK * 10; + auto const stream = std::make_shared(); + tr::SamplingConfig const samplingConfig{kPE_BEAM_WIDTH}; + auto kvCacheManager = makePriorityEvictionManager(blocksInPrimaryPool, maxAttentionWindow, stream); + + auto const& blockManager = kvCacheManager->getBlockManager(); + EXPECT_EQ(kvCacheManager->getNumFreeBlocks(), blocksInPrimaryPool); + + // Seq A: 12 tokens in 3 blocks with strictly ordered priorities. + // B0=[0..3]: MIN priority (0) — will be evicted first (interior node, parent of B1) + // B1=[4..7]: DEFAULT priority — will be evicted second (interior node, parent of B2) + // B2=[8..11]: HIGH priority (90) — will be evicted last (leaf node) + // Trie chain: root → B0 → B1 → B2 + auto inputTokensA = std::make_shared(VecTokens{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}); + auto const inputLengthA = static_cast(inputTokensA->size()); + auto llmRequestA + = std::make_shared(0, kPE_MAX_NEW_TOKENS, inputTokensA, samplingConfig, kPE_IS_STREAMING); + llmRequestA->setKvCacheRetentionConfig(KvCacheRetentionConfig( + {KvCacheRetentionConfig::TokenRangeRetentionConfig(0, 4, KvCacheRetentionConfig::kMinRetentionPriority), + KvCacheRetentionConfig::TokenRangeRetentionConfig(4, 8, KvCacheRetentionConfig::kDefaultRetentionPriority), + KvCacheRetentionConfig::TokenRangeRetentionConfig(8, 12, 90)}, + KvCacheRetentionConfig::kDefaultRetentionPriority)); + kvCacheManager->addSequence(0, inputLengthA, kPE_BEAM_WIDTH, llmRequestA); + kvCacheManager->storeContextBlocks(*llmRequestA); + (void) kvCacheManager->removeSequence(0, llmRequestA); + + // Free queue after release (6 blocks): + // prio 0 (MIN): [B0] ← interior, lowest priority + // prio 35 (DEFAULT): [B3, B4, B5, B1] ← B3/B4/B5 never used; B1 interior + // prio 90 (HIGH): [B2] ← leaf, highest priority + EXPECT_EQ(kvCacheManager->getNumFreeBlocks(), blocksInPrimaryPool); + EXPECT_TRUE(blockManager.verifyQueueIntegrity(maxAttentionWindow)); + + // Step 1: claim 1 block — must take B0 (prio 0, lowest). + // B0 is an interior node (parent of B1→B2 in the trie). + // True priority eviction detaches ONLY B0; B1 and B2 remain in trie. + auto inputTokensX = std::make_shared(VecTokens{200, 201, 202, 203}); + auto const inputLengthX = static_cast(inputTokensX->size()); + auto llmRequestX + = std::make_shared(1, kPE_MAX_NEW_TOKENS, inputTokensX, samplingConfig, kPE_IS_STREAMING); + kvCacheManager->addSequence(1, inputLengthX, kPE_BEAM_WIDTH, llmRequestX); + + // 5 blocks remain after B0 is claimed. + EXPECT_EQ(kvCacheManager->getNumFreeBlocks(), 5); + EXPECT_TRUE(blockManager.verifyQueueIntegrity(maxAttentionWindow)); + + (void) kvCacheManager->removeSequence(1, llmRequestX); + + // Step 2: claim 3 more blocks (all DEFAULT-priority: B3, B4, B5 or B1 depending on queue). + // With true priority eviction, B3, B4, B5 (initialized at DEFAULT ahead of B1 in the queue) + // and B1 (also DEFAULT) are all candidates; B2 (HIGH=90) is still protected. + auto inputTokensY + = std::make_shared(VecTokens{300, 301, 302, 303, 304, 305, 306, 307, 308, 309, 310, 311}); + auto const inputLengthY = static_cast(inputTokensY->size()); + auto llmRequestY + = std::make_shared(2, kPE_MAX_NEW_TOKENS, inputTokensY, samplingConfig, kPE_IS_STREAMING); + kvCacheManager->addSequence(2, inputLengthY, kPE_BEAM_WIDTH, llmRequestY); + + // After seq X is released (returns 1 block) and seq Y claims 3: + // free = 6 (all released by X) - 3 (claimed by Y) = 3 + EXPECT_EQ(kvCacheManager->getNumFreeBlocks(), 3); + EXPECT_TRUE(blockManager.verifyQueueIntegrity(maxAttentionWindow)); + + (void) kvCacheManager->removeSequence(2, llmRequestY); + EXPECT_EQ(kvCacheManager->getNumFreeBlocks(), blocksInPrimaryPool); + EXPECT_TRUE(blockManager.verifyQueueIntegrity(maxAttentionWindow)); + + // Explicit reuse assertion: seq Y stored its 3 blocks ([300..311]) in the trie + // on removeSequence. A new request with the same prefix must be able to reuse at + // least one of those blocks, confirming the trie is consistent after interior-block + // eviction. + auto inputTokensZ = std::make_shared(*inputTokensY); + auto llmRequestZ + = std::make_shared(3, kPE_MAX_NEW_TOKENS, inputTokensZ, samplingConfig, kPE_IS_STREAMING); + kvCacheManager->addSequence(3, static_cast(inputTokensZ->size()), kPE_BEAM_WIDTH, llmRequestZ); + EXPECT_GE(llmRequestZ->getContextCurrentPosition(), kPE_TOKENS_PER_BLOCK); + (void) kvCacheManager->removeSequence(3, llmRequestZ); + EXPECT_EQ(kvCacheManager->getNumFreeBlocks(), blocksInPrimaryPool); + EXPECT_TRUE(blockManager.verifyQueueIntegrity(maxAttentionWindow)); +} + +// Verifies that after a sequence stores blocks in the trie and those blocks are evicted +// via interior-block eviction, subsequent sequences can still allocate and store blocks +// correctly (no trie corruption or assertion failures). +TEST_F(KVCacheManagerTest, TruePriorityEvictionNoCrashAfterInteriorEviction) +{ + auto constexpr blocksInPrimaryPool = 8; + auto const maxAttentionWindow = kPE_TOKENS_PER_BLOCK * 10; + auto const stream = std::make_shared(); + tr::SamplingConfig const samplingConfig{kPE_BEAM_WIDTH}; + auto kvCacheManager = makePriorityEvictionManager(blocksInPrimaryPool, maxAttentionWindow, stream); + + auto const& blockManager = kvCacheManager->getBlockManager(); + + // Seq 0: 3 blocks — MIN priority for first block (interior), DEFAULT for the rest. + // Trie: root → B0(MIN) → B1(DEFAULT) → B2(DEFAULT) + auto inputTokens0 = std::make_shared(VecTokens{0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11}); + auto llmRequest0 + = std::make_shared(0, kPE_MAX_NEW_TOKENS, inputTokens0, samplingConfig, kPE_IS_STREAMING); + llmRequest0->setKvCacheRetentionConfig(KvCacheRetentionConfig( + {KvCacheRetentionConfig::TokenRangeRetentionConfig(0, 4, KvCacheRetentionConfig::kMinRetentionPriority)}, + KvCacheRetentionConfig::kDefaultRetentionPriority)); + kvCacheManager->addSequence(0, static_cast(inputTokens0->size()), kPE_BEAM_WIDTH, llmRequest0); + kvCacheManager->storeContextBlocks(*llmRequest0); + (void) kvCacheManager->removeSequence(0, llmRequest0); + + EXPECT_EQ(kvCacheManager->getNumFreeBlocks(), blocksInPrimaryPool); + EXPECT_TRUE(blockManager.verifyQueueIntegrity(maxAttentionWindow)); + + // Seq 1: 8 completely new tokens — forces eviction of B0 (MIN priority, interior node). + // True priority eviction detaches only B0; B1, B2 remain in trie. + auto inputTokens1 + = std::make_shared(VecTokens{100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111}); + auto llmRequest1 + = std::make_shared(1, kPE_MAX_NEW_TOKENS, inputTokens1, samplingConfig, kPE_IS_STREAMING); + kvCacheManager->addSequence(1, static_cast(inputTokens1->size()), kPE_BEAM_WIDTH, llmRequest1); + kvCacheManager->storeContextBlocks(*llmRequest1); + (void) kvCacheManager->removeSequence(1, llmRequest1); + + EXPECT_TRUE(blockManager.verifyQueueIntegrity(maxAttentionWindow)); + + // Seq 2: 4 new tokens (fresh; no overlap with any prior sequence). + auto inputTokens2 = std::make_shared(VecTokens{200, 201, 202, 203}); + auto llmRequest2 + = std::make_shared(2, kPE_MAX_NEW_TOKENS, inputTokens2, samplingConfig, kPE_IS_STREAMING); + EXPECT_NO_THROW( + kvCacheManager->addSequence(2, static_cast(inputTokens2->size()), kPE_BEAM_WIDTH, llmRequest2)); + kvCacheManager->storeContextBlocks(*llmRequest2); + (void) kvCacheManager->removeSequence(2, llmRequest2); + + EXPECT_EQ(kvCacheManager->getNumFreeBlocks(), blocksInPrimaryPool); + EXPECT_TRUE(blockManager.verifyQueueIntegrity(maxAttentionWindow)); + + // Seq 3: reuses the same tokens as seq 1 — verifies that the interior-eviction + // path left the trie in a consistent state for subsequent insertions/lookups. + auto inputTokens3 + = std::make_shared(VecTokens{100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111}); + auto llmRequest3 + = std::make_shared(3, kPE_MAX_NEW_TOKENS, inputTokens3, samplingConfig, kPE_IS_STREAMING); + EXPECT_NO_THROW( + kvCacheManager->addSequence(3, static_cast(inputTokens3->size()), kPE_BEAM_WIDTH, llmRequest3)); + + // Seq 1's blocks are in the trie and should be reused. + EXPECT_GT(llmRequest3->getContextCurrentPosition(), 0); + + (void) kvCacheManager->removeSequence(3, llmRequest3); + EXPECT_TRUE(blockManager.verifyQueueIntegrity(maxAttentionWindow)); +} + TEST(KVCacheManagerHelpersTest, ChopVectorIntoBlocksBasicNoPartial) { using namespace tensorrt_llm::batch_manager::kv_cache_manager; @@ -4339,8 +4577,13 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStreamWindowSize) events = getEvents(kvCacheManager); - // Expecting only 1 event, storeContextBlock is not called for sliding window. - EXPECT_EQ(events.size(), 1); + // Both window managers store context blocks: one event per window. + EXPECT_EQ(events.size(), 2); + + // BlockManager iterates mWindowBlockManagers in ascending key order, so the + // sliding-window event (smaller windowSize) arrives first. + EXPECT_EQ(events.front().windowSize, slidingWindow); + EXPECT_TRUE(std::holds_alternative(events.front().data)); EXPECT_EQ(events.back().windowSize, maxAttentionWindow); EXPECT_TRUE(std::holds_alternative(events.back().data)); @@ -6344,24 +6587,21 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventRemovedOrderedBeforeStore) << ") for the same window. enqueueStoredEvent must flush pending removes before appending the store."; } -// A store event for window W2 must not flush pending remove events for a different window W1. -// Removes for W1 must only be committed when a store for W1 occurs or when flush() is called. +// A store event for window W must not flush pending remove events for a different window W'. +// Removes for W' must only be committed when a store for W' occurs or when flush() is called. // This verifies per-window isolation in the lazy-batching remove event logic. TEST_F(KVCacheManagerTest, KVCacheManagerEventStoreForDifferentWindowDoesNotFlushPendingRemoves) { // Two windows: wFull (non-SWA, equal to maxSequenceLength) and wSWA (SWA, smaller). - // storeContextBlocks skips SWA windows, so it only emits a Stored event for wFull. - // This means wSWA removes are never flushed by the wFull store — they stay buffered - // until flush() at end of iteration. + // storeContextBlocks iterates mWindowBlockManagers in ascending window-size order + // (std::map), so wSWA is processed before wFull. // - // Expected event order: [Removed(wFull), Stored(wFull), Removed(wSWA)] - // Removed(wFull) — flushed by wFull's own storeContextBlocks call - // Stored(wFull) — emitted by storeContextBlocks for wFull - // Removed(wSWA) — only flushed by the iteration-end flush(), AFTER storeContextBlocks + // Expected event order: [Removed(wSWA), Stored(wSWA), Removed(wFull), Stored(wFull)] + // enqueueStoredEvent(wSWA): flushRemovedEvents(wSWA) → Removed(wSWA), then Stored(wSWA) + // enqueueStoredEvent(wFull): flushRemovedEvents(wFull) → Removed(wFull), then Stored(wFull) // - // If isolation were broken (wFull store flushes ALL windows' removes), the order - // would be [Removed(wSWA), Removed(wFull), Stored(wFull)] — Stored(wFull) would - // appear after Removed(wSWA), violating the per-window ordering guarantee. + // Per-window isolation: enqueueStoredEvent(wSWA) must NOT flush wFull's pending removes. + // If isolation were broken, Removed(wFull) would appear before Stored(wSWA) instead of after. auto constexpr numLayers = 2; auto constexpr numHeads = 2; auto constexpr sizePerHead = 16; @@ -6391,10 +6631,9 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStoreForDifferentWindowDoesNotFlus kvCacheManager.allocatePools(false); (void) getEvents(kvCacheManager); - // Seq0: 9 tokens → 3 blocks per window. storeContextBlocks stores 2 full blocks in wFull - // (skips wSWA). removeSequence stores 2 full blocks in wSWA as well (releaseBlocks covers - // all windows). After release, each window's free queue is [block3_fresh, block2, block1, block0], - // with block0 and block1 in the respective radix trees. + // Seq0: 9 tokens → 3 blocks per window. storeContextBlocks stores 2 full blocks in both + // wSWA and wFull. After removeSequence, each window's free queue is + // [block3_fresh, block2, block1, block0], with block0 and block1 in their radix trees. auto inputTokens0 = std::make_shared(VecTokens{0, 1, 2, 3, 4, 5, 6, 7, 8}); auto llmRequest0 = std::make_shared(0, maxNewTokens, inputTokens0, samplingConfig, true); kvCacheManager.addSequence(0, inputTokens0->size(), beamWidth, llmRequest0); @@ -6405,10 +6644,10 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStoreForDifferentWindowDoesNotFlus // Seq1 with different tokens (9 tokens → 3 blocks per window). // addSequence for each window: gets block3 (fresh, no event), block2 (not in tree, no event), // then block1 (in tree as leaf) → freeChildren(block1) → Removed(block1) buffered for that window. - // storeContextBlocks: - // wSWA: skipped (SWA) — wSWA removes stay buffered - // wFull: flushRemovedEvents(wFull) → Removed(wFull) committed; Stored(wFull) committed - // flush(): flushRemovedEvents(wSWA) → Removed(wSWA) committed + // storeContextBlocks (ascending window-size order): + // wSWA first: flushRemovedEvents(wSWA) → Removed(wSWA) committed; Stored(wSWA) committed + // wFull next: flushRemovedEvents(wFull) → Removed(wFull) committed; Stored(wFull) committed + // The wSWA store must NOT have flushed wFull's pending removes prematurely. auto inputTokens1 = std::make_shared(VecTokens{100, 101, 102, 103, 104, 105, 106, 107, 108}); auto llmRequest1 = std::make_shared(1, maxNewTokens, inputTokens1, samplingConfig, true); kvCacheManager.addSequence(1, inputTokens1->size(), beamWidth, llmRequest1); @@ -6417,7 +6656,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStoreForDifferentWindowDoesNotFlus auto events = getEvents(kvCacheManager); // Find the position of the first Removed and Stored event for each window. - std::optional removedSWAPos, storedFullPos, removedFullPos; + std::optional removedSWAPos, storedSWAPos, removedFullPos, storedFullPos; SizeType32 pos = 0; for (auto const& event : events) { @@ -6430,25 +6669,752 @@ TEST_F(KVCacheManagerTest, KVCacheManagerEventStoreForDifferentWindowDoesNotFlus } else if (std::holds_alternative(event.data)) { + if (event.windowSize == wSWA && !storedSWAPos) + storedSWAPos = pos; if (event.windowSize == wFull && !storedFullPos) - { storedFullPos = pos; - } } ++pos; } ASSERT_TRUE(removedSWAPos.has_value()) << "Expected Removed event for wSWA"; + ASSERT_TRUE(storedSWAPos.has_value()) << "Expected Stored event for wSWA"; ASSERT_TRUE(removedFullPos.has_value()) << "Expected Removed event for wFull"; ASSERT_TRUE(storedFullPos.has_value()) << "Expected Stored event for wFull"; - // Within wFull, removes must precede stores. + // Within each window, removes must precede stores (per-window ordering guarantee). + EXPECT_LT(*removedSWAPos, *storedSWAPos) << "Removed(wSWA) must precede Stored(wSWA)"; EXPECT_LT(*removedFullPos, *storedFullPos) << "Removed(wFull) must precede Stored(wFull)"; - // The wFull store must NOT have flushed wSWA's pending removes prematurely. - // Correct isolation: Stored(wFull) appears before Removed(wSWA). - // Broken isolation: Removed(wSWA) appears before Stored(wFull). - EXPECT_LT(*storedFullPos, *removedSWAPos) - << "Stored(wFull) (pos=" << *storedFullPos << ") must precede Removed(wSWA) (pos=" << *removedSWAPos - << "). The wFull store must not prematurely flush pending removes for wSWA."; + // The wSWA store must NOT have flushed wFull's pending removes prematurely. + // Correct isolation: Stored(wSWA) appears before Removed(wFull). + // Broken isolation: Removed(wFull) appears before Stored(wSWA). + EXPECT_LT(*storedSWAPos, *removedFullPos) + << "Stored(wSWA) (pos=" << *storedSWAPos << ") must precede Removed(wFull) (pos=" << *removedFullPos + << "). The wSWA store must not prematurely flush pending removes for wFull."; +} + +// --------------------------------------------------------------------------- +// VSWA tests +// --------------------------------------------------------------------------- + +namespace +{ +// Shared constants for all VSWA tests. +auto constexpr kVSWA_TOKENS_PER_BLOCK = 4; +auto constexpr kVSWA_ATTENTION_WINDOW = 8; +auto constexpr kVSWA_MAX_SEQUENCE_LENGTH = 128; +SizeType32 constexpr kVSWA_MAX_NEW_TOKENS = 40; +auto constexpr kVSWA_BEAM_WIDTH = 1; +auto constexpr kVSWA_BEAM_IDX = 0; +bool constexpr kVSWA_IS_STREAMING = false; +TokenIdType constexpr kVSWA_FIRST_TOKEN = 1000; + +// Factory: construct and allocate a KVCacheManager for VSWA tests. +// numLayers=2, numHeads=2, sizePerHead=64, tokensPerBlock=4, attentionWindow=8, +// maxNumSequences=8, beamWidth=1, sinkTokenLength=0, maxSequenceLength=128. +std::unique_ptr makeVSWAManager( + SizeType32 blocksInPrimaryPool, bool enableBlockReuse, std::shared_ptr const& stream) +{ + auto const blocksPerWindow = BlocksPerWindow{{kVSWA_ATTENTION_WINDOW, {blocksInPrimaryPool, 0}}}; + auto mgr = std::make_unique(2, 2, 64, kVSWA_TOKENS_PER_BLOCK, blocksPerWindow, 8, kVSWA_BEAM_WIDTH, + std::vector{kVSWA_ATTENTION_WINDOW}, std::nullopt, nvinfer1::DataType::kHALF, 0, stream, + kVSWA_MAX_SEQUENCE_LENGTH, enableBlockReuse, /*onboardBlocks=*/false); + mgr->allocatePools(false); + return mgr; +} + +// Factory: construct and allocate a KVCacheManager with window==tokensPerBlock for +// multi-OOW tests. With window=4 and tpb=4 the OOW condition fires at numTokens=8, +// so two consecutive addToken calls (after 11 context tokens) cause two OOW events +// before the next block boundary — exercising the prevBlock->isPlaceholder() path in +// storeNewBlock. +std::unique_ptr makeSmallWindowManager( + SizeType32 blocksInPrimaryPool, std::shared_ptr const& stream) +{ + SizeType32 constexpr kSmallWindow = 4; + SizeType32 constexpr kSmallTpb = 4; + SizeType32 constexpr kSmallMaxSeqLen = 128; + auto const blocksPerWindow = BlocksPerWindow{{kSmallWindow, {blocksInPrimaryPool, 0}}}; + auto mgr = std::make_unique(2, 2, 64, kSmallTpb, blocksPerWindow, 8, kVSWA_BEAM_WIDTH, + std::vector{kSmallWindow}, std::nullopt, nvinfer1::DataType::kHALF, 0, stream, kSmallMaxSeqLen, + /*enableBlockReuse=*/true, /*onboardBlocks=*/false); + mgr->allocatePools(false); + return mgr; +} +} // namespace + +// Verify that a non-stolen OOW block (hasRefs() == 0 at releaseBlocks time) is +// stored in the reuse trie and can be reused by a subsequent sequence. +TEST_F(KVCacheManagerTest, VSWANonStolenOOWBlockStoredForReuse) +{ + // SWA with a generous pool so the OOW block is never stolen. + auto constexpr blocksInPrimaryPool = 8; + auto const stream = std::make_shared(); + tr::SamplingConfig const samplingConfig{kVSWA_BEAM_WIDTH}; + auto kvCacheManager = makeVSWAManager(blocksInPrimaryPool, /*enableBlockReuse=*/true, stream); + + auto const& blockManager = kvCacheManager->getBlockManager(); + TokenIdType constexpr firstToken = kVSWA_FIRST_TOKEN; + + // Seq 0: 11 input tokens → allocates 3 blocks covering tokens [1000..1010]. + // After addToken (token 1011), numTokens==12 triggers OOW for block 0 (tokens + // [1000..1003]) which enters the free queue at MIN priority with hasRefs()==0. + auto inputTokens0 = std::make_shared(11); + std::iota(inputTokens0->begin(), inputTokens0->end(), firstToken); + auto llmRequest0 + = std::make_shared(0, kVSWA_MAX_NEW_TOKENS, inputTokens0, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(0, 11, kVSWA_BEAM_WIDTH, llmRequest0); + // Store B0 and B1 in the reuse trie so they are there before B0 goes OOW. + kvCacheManager->storeContextBlocks(*llmRequest0); + + llmRequest0->addNewToken(firstToken + 11, kVSWA_BEAM_IDX); + kvCacheManager->addToken(0); + + // Release seq 0: the placeholder at position 0 → storeBlocks sees node K0 still has value B0 + // (not stolen) → advances prevBlock → B0's chain stored for reuse. + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(0, llmRequest0))); + + // Seq 1: exactly the same 4-token prefix as the OOW block → must reuse it. + auto inputTokens1 = std::make_shared(kVSWA_TOKENS_PER_BLOCK); + std::iota(inputTokens1->begin(), inputTokens1->end(), firstToken); + auto llmRequest1 + = std::make_shared(1, kVSWA_MAX_NEW_TOKENS, inputTokens1, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(1, kVSWA_TOKENS_PER_BLOCK, kVSWA_BEAM_WIDTH, llmRequest1); + + // The OOW block was stored with 4 tokens, but S1's usableSize=4-1=3 so the + // search key has 3 tokens. 3/4 tokens match → contextCurrentPosition == 3. + // Any non-zero value confirms the OOW block was stored and is being reused. + EXPECT_EQ(llmRequest1->getContextCurrentPosition(), kVSWA_TOKENS_PER_BLOCK - 1); + + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(1, llmRequest1))); + // All blocks must be free — no leaks. + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); +} + +// Verify that storeNewBlock stores SWA blocks (including OOW blocks) into the reuse +// trie during generation — i.e., without waiting for removeSequence. +// After two generation steps (reaching a block boundary at usableSize=12), blocks +// B0 (OOW), B1, and B2 are stored. A subsequent sequence can then reuse B0 while +// seq0 is still alive. +TEST_F(KVCacheManagerTest, VSWABlockStoredDuringGeneration) +{ + // Generous pool so no blocks are stolen. + auto constexpr blocksInPrimaryPool = 10; + auto const stream = std::make_shared(); + tr::SamplingConfig const samplingConfig{kVSWA_BEAM_WIDTH}; + auto kvCacheManager = makeVSWAManager(blocksInPrimaryPool, /*enableBlockReuse=*/true, stream); + auto const& blockManager = kvCacheManager->getBlockManager(); + + // Seq 0: 11 input tokens covering blocks B0=[1000..1003], B1=[1004..1007], B2=[1008..1010] (partial). + auto inputTokens0 = std::make_shared(11); + std::iota(inputTokens0->begin(), inputTokens0->end(), kVSWA_FIRST_TOKEN); + auto llmRequest0 + = std::make_shared(0, kVSWA_MAX_NEW_TOKENS, inputTokens0, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(0, 11, kVSWA_BEAM_WIDTH, llmRequest0); + // Store B0 and B1 in the reuse trie during context (invariant: stored before OOW). + kvCacheManager->storeContextBlocks(*llmRequest0); + + // Generation step 1: token 1011. + // numTokens becomes 12; usableSize=11, 11%4!=0 → storeNewBlock is a no-op. + // adjustBlocksIfNeeded: 12-0*4=12 >= 8+4=12 → B0 goes OOW (detachFrontBlock). + llmRequest0->addNewToken(kVSWA_FIRST_TOKEN + 11, kVSWA_BEAM_IDX); + kvCacheManager->addToken(0); + kvCacheManager->storeNewBlock(*llmRequest0); // no-op (usableSize=11) + + // Generation step 2: token 1012. + // numTokens becomes 13; usableSize=12, 12%4==0 → storeNewBlock fires. + // storeNewBlock processes [P0, B1, B2]: P0→node K0 has value B0→advance; + // B1→node K1 has value B1 (from context)→advance; B2→node K2 empty→insert. + // adjustBlocksIfNeeded: 13-1*4=9 < 12 → no additional OOW detach. + // (13-1)%4==0 → a new block B3 is allocated for position 3. + llmRequest0->addNewToken(kVSWA_FIRST_TOKEN + 12, kVSWA_BEAM_IDX); + kvCacheManager->addToken(0); + kvCacheManager->storeNewBlock(*llmRequest0); // stores B2 (B0+B1 already in trie from context) + + // Seq 1: same 4-token prefix as B0 → should reuse it without seq0 being released. + auto inputTokens1 = std::make_shared(kVSWA_TOKENS_PER_BLOCK); + std::iota(inputTokens1->begin(), inputTokens1->end(), kVSWA_FIRST_TOKEN); + auto llmRequest1 + = std::make_shared(1, kVSWA_MAX_NEW_TOKENS, inputTokens1, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(1, kVSWA_TOKENS_PER_BLOCK, kVSWA_BEAM_WIDTH, llmRequest1); + + // B0 was stored during generation (not just at release time). + // usableSize for seq1 context = 4-1=3 tokens → partial match of 3 tokens. + EXPECT_EQ(llmRequest1->getContextCurrentPosition(), kVSWA_TOKENS_PER_BLOCK - 1); + + // Clean up both sequences. + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(1, llmRequest1))); + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(0, std::nullopt))); + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); +} + +// Verify that when an OOW block is stolen by another sequence, storeBlocks stops +// at that block (hasRefs() > 0) without corrupting the acquiring sequence's trie, +// and all blocks are properly released on removeSequence for both sequences. +TEST_F(KVCacheManagerTest, VSWAStolenOOWBlockNoCorruption) +{ + // Tight pool: seq0 needs 3 context blocks + 1 for addToken = 4 total. + // Seq1 needs 2 blocks. The one block in the free queue after seq0's addToken + // goes to seq1, which steals the OOW block. + auto constexpr blocksInPrimaryPool = 4; + auto const stream = std::make_shared(); + tr::SamplingConfig const samplingConfig{kVSWA_BEAM_WIDTH}; + auto kvCacheManager = makeVSWAManager(blocksInPrimaryPool, /*enableBlockReuse=*/true, stream); + + auto const& blockManager = kvCacheManager->getBlockManager(); + + // Seq 0: 11 tokens, triggering 1 OOW block after addToken. + auto inputTokens0 = std::make_shared(11); + std::iota(inputTokens0->begin(), inputTokens0->end(), kVSWA_FIRST_TOKEN); + auto llmRequest0 + = std::make_shared(0, kVSWA_MAX_NEW_TOKENS, inputTokens0, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(0, 11, kVSWA_BEAM_WIDTH, llmRequest0); + // Store B0 and B1 in the trie before they can go OOW. + kvCacheManager->storeContextBlocks(*llmRequest0); + + llmRequest0->addNewToken(kVSWA_FIRST_TOKEN + 11, kVSWA_BEAM_IDX); + kvCacheManager->addToken(0); + + // After addToken: B0 goes OOW (detachFrontBlock); (12-1)%4 != 0 so no new + // block is allocated. S0 holds B1, B2 in-window. Pool=4: B0(DEFAULT) + B3(DEFAULT) + // = 2 free blocks. B0 is still in the trie (stored by storeContextBlocks). + EXPECT_EQ(blockManager.getNumFreeBlocks(), 2); + + // Seq 1: 8 tokens → needs 2 blocks. It acquires B3 (DEFAULT, oldest) and B0 (DEFAULT), + // stealing the OOW block away from seq 0. getFreeBlock(B0) calls detachFromLookupNode, + // removing B0 from the trie. + auto inputTokens1 = std::make_shared(8); + std::iota(inputTokens1->begin(), inputTokens1->end(), kVSWA_FIRST_TOKEN + 100); + auto llmRequest1 + = std::make_shared(1, kVSWA_MAX_NEW_TOKENS, inputTokens1, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(1, 8, kVSWA_BEAM_WIDTH, llmRequest1); + + // Seq 0's removeSequence: storeBlocks sees placeholder P0 → node K0 has no value + // (B0 was detached from trie when seq1's getFreeBlock claimed it) → stops cleanly. + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(0, llmRequest0))); + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(1, llmRequest1))); + + // All blocks must be free after both sequences are released. + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); + + // Reuse assertion: seq1 stored its 2 blocks ([kVSWA_FIRST_TOKEN+100 .. +107]) in the + // trie during removeSequence. A follow-up request with seq1's prefix must be able to + // reuse at least one of those blocks, confirming that seq0's storeBlocks correctly + // stopped at the stolen OOW block and did NOT corrupt the trie with seq0's stale prefix. + auto inputTokensReuse = std::make_shared(*inputTokens1); + auto llmRequestReuse + = std::make_shared(2, kVSWA_MAX_NEW_TOKENS, inputTokensReuse, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence( + 2, static_cast(inputTokensReuse->size()), kVSWA_BEAM_WIDTH, llmRequestReuse); + EXPECT_GT(llmRequestReuse->getContextCurrentPosition(), 0); + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(2, llmRequestReuse))); + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); +} + +// Verify the placeholder path when the acquiring sequence finishes (removeSequence) BEFORE +// the original sequence: the OOW block has hasRefs()==false but is stored in the trie +// under the acquirer's key. storeBlocks for the original sequence encounters a placeholder +// at the OOW position; the trie node for K_seq0_block0 has no value (block stored at +// seq1's key, not seq0's) → breaks, preserving the acquirer's trie entry for reuse. +TEST_F(KVCacheManagerTest, VSWAStolenAndReleasedOOWBlockIsInLookupTreeProtection) +{ + // Pool=3: seq0 uses all 3 blocks (B0..B2) for context. After addToken, only B0 is + // in the free queue (no B3 exists), so seq1 must take B0 — the stolen OOW block. + auto constexpr blocksInPrimaryPool = 3; + auto const stream = std::make_shared(); + tr::SamplingConfig const samplingConfig{kVSWA_BEAM_WIDTH}; + auto kvCacheManager = makeVSWAManager(blocksInPrimaryPool, /*enableBlockReuse=*/true, stream); + + auto const& blockManager = kvCacheManager->getBlockManager(); + TokenIdType constexpr seq1FirstToken = 1100; + + // Seq 0: 11 tokens → allocates B0 (tokens 1000..1003), B1 (1004..1007), B2 (1008..1010). + auto inputTokens0 = std::make_shared(11); + std::iota(inputTokens0->begin(), inputTokens0->end(), kVSWA_FIRST_TOKEN); + auto llmRequest0 + = std::make_shared(0, kVSWA_MAX_NEW_TOKENS, inputTokens0, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(0, 11, kVSWA_BEAM_WIDTH, llmRequest0); + // Store B0 and B1 in the trie before B0 goes OOW. + kvCacheManager->storeContextBlocks(*llmRequest0); + + // addToken: B0 goes OOW at DEFAULT priority; no new block allocated ((12-1)%4 != 0). + // Free queue: [B0] — the only free block in the pool. + llmRequest0->addNewToken(kVSWA_FIRST_TOKEN + 11, kVSWA_BEAM_IDX); + kvCacheManager->addToken(0); + EXPECT_EQ(blockManager.getNumFreeBlocks(), 1); + + // Seq 1: 4 tokens (distinct prefix) → steals B0 (the only free block). + // getFreeBlock(B0) calls detachFromLookupNode, removing B0 from the trie. + auto inputTokens1 = std::make_shared(kVSWA_TOKENS_PER_BLOCK); + std::iota(inputTokens1->begin(), inputTokens1->end(), seq1FirstToken); + auto llmRequest1 + = std::make_shared(1, kVSWA_MAX_NEW_TOKENS, inputTokens1, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(1, kVSWA_TOKENS_PER_BLOCK, kVSWA_BEAM_WIDTH, llmRequest1); + EXPECT_EQ(blockManager.getNumFreeBlocks(), 0); // pool exhausted + + // removeSequence(1) FIRST: seq1's storeBlocks stores B0 (now holding seq1's tokens) in + // the trie under seq1's key. B0 is no longer in the trie at seq0's key. + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(1, llmRequest1))); + EXPECT_EQ(blockManager.getNumFreeBlocks(), 1); // B0 freed into free queue + + // removeSequence(0): seq0's storeBlocks encounters P0 (placeholder) at position 0. + // The trie node for K_seq0_block0 has no value (B0 is stored at seq1's key, not seq0's). + // Placeholder path: anchor evicted → break immediately. No crash, no trie corruption. + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(0, llmRequest0))); + + // All 3 blocks must be free (no leaks). + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); + + // Seq 2: same prefix as seq1 → must reuse B0 stored at seq1's key. + auto inputTokens2 = std::make_shared(kVSWA_TOKENS_PER_BLOCK); + std::iota(inputTokens2->begin(), inputTokens2->end(), seq1FirstToken); + auto llmRequest2 + = std::make_shared(2, kVSWA_MAX_NEW_TOKENS, inputTokens2, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(2, kVSWA_TOKENS_PER_BLOCK, kVSWA_BEAM_WIDTH, llmRequest2); + // 4 tokens stored, usable key has 3 tokens → 3/4 match → contextCurrentPosition==3. + EXPECT_EQ(llmRequest2->getContextCurrentPosition(), kVSWA_TOKENS_PER_BLOCK - 1); + + // Seq 3: same prefix as seq0's first block → must NOT find it (chain broke at placeholder P0, + // so seq0's blocks were never stored; B0 is only in the trie at seq1's key). + auto inputTokens3 = std::make_shared(kVSWA_TOKENS_PER_BLOCK); + std::iota(inputTokens3->begin(), inputTokens3->end(), kVSWA_FIRST_TOKEN); + auto llmRequest3 + = std::make_shared(3, kVSWA_MAX_NEW_TOKENS, inputTokens3, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(3, kVSWA_TOKENS_PER_BLOCK, kVSWA_BEAM_WIDTH, llmRequest3); + EXPECT_EQ(llmRequest3->getContextCurrentPosition(), 0); + + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(2, llmRequest2))); + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(3, llmRequest3))); + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); +} + +// Verify that OOW blocks are released at their original DEFAULT priority in detachFrontBlock, +// making them the first candidates for eviction over untouched DEFAULT-priority blocks. +TEST_F(KVCacheManagerTest, VSWAOOWBlockReleasedAtOriginalPriority) +{ + // Pool of 5 blocks: seq0 uses B0,B1,B2 (context) + allocates B3 on the block + // boundary after addToken. B4 is the single untouched DEFAULT-priority free block. + // The OOW block B0 enters the free queue at its original DEFAULT priority — it is + // NOT forced to MIN priority. The next allocation follows normal LRU order among + // equal-priority blocks and does NOT preferentially claim B0. + auto constexpr blocksInPrimaryPool = 5; + auto const stream = std::make_shared(); + tr::SamplingConfig const samplingConfig{kVSWA_BEAM_WIDTH}; + // Use reuse=false so seq1's addSequence does a plain allocation (no trie lookup). + auto kvCacheManager = makeVSWAManager(blocksInPrimaryPool, /*enableBlockReuse=*/false, stream); + + auto const& blockManager = kvCacheManager->getBlockManager(); + + // Seq 0: 11 input tokens → allocates B0, B1, B2. + auto inputTokens0 = std::make_shared(11); + std::iota(inputTokens0->begin(), inputTokens0->end(), kVSWA_FIRST_TOKEN); + auto llmRequest0 + = std::make_shared(0, kVSWA_MAX_NEW_TOKENS, inputTokens0, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(0, 11, kVSWA_BEAM_WIDTH, llmRequest0); + + // Capture B0's ID before it goes OOW. + auto const onlyWindowSize = theOnlyWindowSize(*kvCacheManager); + auto const& seq0 = kvCacheManager->getSequence(0); + auto const oowBlockId = seq0.getCacheBlockIds(onlyWindowSize)[kVSWA_BEAM_IDX][0]; + + // addToken: B0 → free queue at DEFAULT (original) priority; B3 allocated (block boundary). + // Free queue now: B0 (DEFAULT), B4 (DEFAULT) — same priority, LRU order applies. + llmRequest0->addNewToken(kVSWA_FIRST_TOKEN + 11, kVSWA_BEAM_IDX); + kvCacheManager->addToken(0); + + // After addToken: B0 OOW (DEFAULT priority). (12-1)%4 != 0 → no new block allocated. + // S0 holds B1, B2. Pool=5: B0(DEFAULT) + B3(DEFAULT) + B4(DEFAULT) = 3 free blocks. + EXPECT_EQ(blockManager.getNumFreeBlocks(), 3); + + // Seq 1: 4 tokens → needs 1 block. Both B0 and B4 have DEFAULT priority; LRU picks + // the block that has been free the longest (B4 was never used), NOT B0 (just released). + auto inputTokens1 = std::make_shared(kVSWA_TOKENS_PER_BLOCK); + std::iota(inputTokens1->begin(), inputTokens1->end(), 2000); + auto llmRequest1 + = std::make_shared(1, kVSWA_MAX_NEW_TOKENS, inputTokens1, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(1, kVSWA_TOKENS_PER_BLOCK, kVSWA_BEAM_WIDTH, llmRequest1); + + auto const& seq1 = kvCacheManager->getSequence(1); + auto const seq1BlockId = seq1.getCacheBlockIds(onlyWindowSize)[kVSWA_BEAM_IDX][0]; + + // OOW block (B0, DEFAULT priority) must NOT have been preferentially chosen over + // the other free blocks — its priority is unchanged from context time. + EXPECT_NE(seq1BlockId, oowBlockId); + + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(0, llmRequest0))); + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(1, llmRequest1))); + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); +} + +// Verify the placeholder approach: when a stolen OOW block is freed by its new owner +// WITHOUT being stored in the reuse trie (e.g., the new owner's removeSequence is called +// with std::nullopt, simulating a failed/cancelled request), storeBlocks for the original +// sequence encounters a placeholder at the OOW position, finds no trie entry (anchor +// block was evicted), and stops — preserving trie correctness. +TEST_F(KVCacheManagerTest, VSWAStolenOOWBlockPlaceholderStopsChainStore) +{ + // Pool=3: seq0 uses all 3 blocks (B0..B2) for context. After addToken, only B0 is + // in the free queue, so seq1 (1 block) must take B0 — the stolen OOW block. + auto constexpr blocksInPrimaryPool = 3; + auto const stream = std::make_shared(); + tr::SamplingConfig const samplingConfig{kVSWA_BEAM_WIDTH}; + auto kvCacheManager = makeVSWAManager(blocksInPrimaryPool, /*enableBlockReuse=*/true, stream); + + auto const& blockManager = kvCacheManager->getBlockManager(); + + // Seq 0: 11 tokens → allocates B0 (tokens 1000..1003), B1 (1004..1007), B2 (1008..1010). + // storeContextBlocks stores B0 and B1 in the trie (B2 is partial, not stored). + auto inputTokens0 = std::make_shared(11); + std::iota(inputTokens0->begin(), inputTokens0->end(), kVSWA_FIRST_TOKEN); + auto llmRequest0 + = std::make_shared(0, kVSWA_MAX_NEW_TOKENS, inputTokens0, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(0, 11, kVSWA_BEAM_WIDTH, llmRequest0); + // Store B0 and B1 in the trie before B0 goes OOW. + kvCacheManager->storeContextBlocks(*llmRequest0); + + // addToken: B0 goes OOW at DEFAULT priority → placeholder P0 replaces B0 in seq0's block list. + // (12-1)%4 != 0 → no new block. Free queue: [B0] — the only free block in the pool. + llmRequest0->addNewToken(kVSWA_FIRST_TOKEN + 11, kVSWA_BEAM_IDX); + kvCacheManager->addToken(0); + EXPECT_EQ(blockManager.getNumFreeBlocks(), 1); + + // Seq 1: 4 tokens (distinct prefix, starting at 2000) → steals B0 (the only free block). + // getFreeBlock calls detachFromLookupNode(B0), removing B0 from the trie. + TokenIdType constexpr seq1FirstToken = 2000; + auto inputTokens1 = std::make_shared(kVSWA_TOKENS_PER_BLOCK); + std::iota(inputTokens1->begin(), inputTokens1->end(), seq1FirstToken); + auto llmRequest1 + = std::make_shared(1, kVSWA_MAX_NEW_TOKENS, inputTokens1, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(1, kVSWA_TOKENS_PER_BLOCK, kVSWA_BEAM_WIDTH, llmRequest1); + EXPECT_EQ(blockManager.getNumFreeBlocks(), 0); // pool exhausted + + // Release seq1 with std::nullopt: simulates a failed/cancelled request. + // B0 is freed back to the pool without being stored — it is no longer in the trie. + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(1, std::nullopt))); + EXPECT_EQ(blockManager.getNumFreeBlocks(), 1); // B0 freed + + // Release seq0: storeBlocks encounters P0 (placeholder) at position 0. + // The trie node for B0's original key has no value — B0 was removed from the trie + // by seq1's getFreeBlock → storeBlocks breaks immediately. B1 and B2 are NOT stored + // (the chain is stopped at P0). + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(0, llmRequest0))); + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); + + // Negative reuse check for B0: seq2 with seq0's OOW block prefix must NOT find it — + // the placeholder stopped storeBlocks before B0 could be incorrectly stored under seq0's key. + auto inputTokens2 = std::make_shared(kVSWA_TOKENS_PER_BLOCK); + std::iota(inputTokens2->begin(), inputTokens2->end(), kVSWA_FIRST_TOKEN); + auto llmRequest2 + = std::make_shared(2, kVSWA_MAX_NEW_TOKENS, inputTokens2, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(2, kVSWA_TOKENS_PER_BLOCK, kVSWA_BEAM_WIDTH, llmRequest2); + EXPECT_EQ(llmRequest2->getContextCurrentPosition(), 0); + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(2, llmRequest2))); + + // Chain-stop check for B1: the chain broke at P0, so B1 must also not be in the trie. + auto inputTokens3 = std::make_shared(kVSWA_TOKENS_PER_BLOCK); + std::iota(inputTokens3->begin(), inputTokens3->end(), kVSWA_FIRST_TOKEN + kVSWA_TOKENS_PER_BLOCK); + auto llmRequest3 + = std::make_shared(3, kVSWA_MAX_NEW_TOKENS, inputTokens3, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(3, kVSWA_TOKENS_PER_BLOCK, kVSWA_BEAM_WIDTH, llmRequest3); + EXPECT_EQ(llmRequest3->getContextCurrentPosition(), 0); + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(3, llmRequest3))); + + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); +} + +// Verify that detachFrontBlock replaces the OOW block slot with a placeholder, and that +// when the OOW block is still in the trie, storeBlocks advances the search root past the +// placeholder (chain preserved) so subsequent in-window blocks are stored correctly. +TEST_F(KVCacheManagerTest, VSWAPlaceholderAdvancesSearchRootWhenOOWBlockInTrie) +{ + // Generous pool: no blocks stolen. + auto constexpr blocksInPrimaryPool = 8; + auto const stream = std::make_shared(); + tr::SamplingConfig const samplingConfig{kVSWA_BEAM_WIDTH}; + auto kvCacheManager = makeVSWAManager(blocksInPrimaryPool, /*enableBlockReuse=*/true, stream); + auto const& blockManager = kvCacheManager->getBlockManager(); + + // Seq 0: 11 tokens → B0=[1000..1003], B1=[1004..1007], B2=[1008..1010]. + // storeContextBlocks stores B0 and B1 during context. + auto inputTokens0 = std::make_shared(11); + std::iota(inputTokens0->begin(), inputTokens0->end(), kVSWA_FIRST_TOKEN); + auto llmRequest0 + = std::make_shared(0, kVSWA_MAX_NEW_TOKENS, inputTokens0, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(0, 11, kVSWA_BEAM_WIDTH, llmRequest0); + // Store B0 and B1 in the reuse trie during context (invariant: stored before OOW). + kvCacheManager->storeContextBlocks(*llmRequest0); + + // addToken step 1 (token 1011): B0 goes OOW → P0 placeholder replaces B0 in block list. + // B0 remains in the trie at DEFAULT priority (it was stored by storeContextBlocks). + llmRequest0->addNewToken(kVSWA_FIRST_TOKEN + 11, kVSWA_BEAM_IDX); + kvCacheManager->addToken(0); + kvCacheManager->storeNewBlock(*llmRequest0); // usableSize=11, no-op + + // addToken step 2 (token 1012): (13-1)%4=0 → block boundary, B3 allocated. + // storeNewBlock fires with usableSize=12: processes [P0, B1, B2]. + // insertNodes([K0,K1,K2]) finds/creates all nodes. + // P0 (placeholder) → node K0 has value B0 (still in trie) → advance prevBlock. + // B1 → node K1 has value B1 (from context) → slot occupied → advance prevBlock. + // B2 → node K2 is empty → insert B2 into trie. + llmRequest0->addNewToken(kVSWA_FIRST_TOKEN + 12, kVSWA_BEAM_IDX); + kvCacheManager->addToken(0); + kvCacheManager->storeNewBlock(*llmRequest0); // stores B2 + + // Seq 1: same prefix as B0 ([1000..1003]) → must reuse B0 from the trie. + auto inputTokens1 = std::make_shared(kVSWA_TOKENS_PER_BLOCK); + std::iota(inputTokens1->begin(), inputTokens1->end(), kVSWA_FIRST_TOKEN); + auto llmRequest1 + = std::make_shared(1, kVSWA_MAX_NEW_TOKENS, inputTokens1, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(1, kVSWA_TOKENS_PER_BLOCK, kVSWA_BEAM_WIDTH, llmRequest1); + // B0 stored with 4 tokens; seq1's usable key has 3 tokens → 3/4 match. + EXPECT_EQ(llmRequest1->getContextCurrentPosition(), kVSWA_TOKENS_PER_BLOCK - 1); + + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(1, llmRequest1))); + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(0, std::nullopt))); + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); +} + +// Verify that schedulingRemoveSequence correctly skips placeholder blocks and does NOT +// call decSchedulingRefCount() on them (which would trigger a TLLM_CHECK since the +// scheduling ref count of a freshly-created placeholder is 0). +// +// After storeContextBlocks + addToken: +// seq0.mAllocatedBlocksPerSeq = [P0 (placeholder), B1, B2] +// startScheduling() copies mRefCount → mSchedulingRefCount for every slot, including P0 +// (mSchedulingRefCount = 0 for a placeholder). +// schedulingRemoveSequence must skip P0 and only decrement B1 and B2, making all +// blocksInPrimaryPool available from the scheduler's perspective. +TEST_F(KVCacheManagerTest, VSWASchedulingRemoveSequenceSkipsPlaceholders) +{ + // Pool=5: seq0 uses B0, B1, B2 for context; B3, B4 are free throughout. + auto constexpr blocksInPrimaryPool = 5; + auto const stream = std::make_shared(); + tr::SamplingConfig const samplingConfig{kVSWA_BEAM_WIDTH}; + auto kvCacheManager = makeVSWAManager(blocksInPrimaryPool, /*enableBlockReuse=*/true, stream); + auto const& blockManager = kvCacheManager->getBlockManager(); + + // Seq 0: 11 input tokens → B0=[1000..1003], B1=[1004..1007], B2=[1008..1010]. + // storeContextBlocks stores B0 and B1 before they can go OOW. + auto inputTokens0 = std::make_shared(11); + std::iota(inputTokens0->begin(), inputTokens0->end(), kVSWA_FIRST_TOKEN); + auto llmRequest0 + = std::make_shared(0, kVSWA_MAX_NEW_TOKENS, inputTokens0, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(0, 11, kVSWA_BEAM_WIDTH, llmRequest0); + kvCacheManager->storeContextBlocks(*llmRequest0); + + // addToken: B0 goes OOW → detachFrontBlock replaces B0 with placeholder P0. + // P0 has mRefCount=0 and isPlaceholder()==true. + // (12-1)%4 != 0 → no new block. Free pool: B0(DEFAULT), B3(DEFAULT), B4(DEFAULT) = 3 free. + llmRequest0->addNewToken(kVSWA_FIRST_TOKEN + 11, kVSWA_BEAM_IDX); + kvCacheManager->addToken(0); + EXPECT_EQ(blockManager.getNumFreeBlocks(), 3); + + // startScheduling() snapshots free blocks and copies mRefCount → mSchedulingRefCount + // for every allocated block, including the placeholder P0 (mSchedulingRefCount = 0). + kvCacheManager->startScheduling(); + + // schedulingRemoveSequence must skip P0 (isPlaceholder()==true) and only decrement + // the scheduling ref counts of B1 and B2. Calling decSchedulingRefCount() on P0 + // (with mSchedulingRefCount=0) would fire TLLM_CHECK_WITH_INFO and abort the test. + EXPECT_NO_THROW(kvCacheManager->schedulingRemoveSequence(0)); + + // After skipping P0 and releasing B1+B2, all 5 blocks are available for scheduling. + EXPECT_TRUE(blockManager.schedulingHasFreeBlocks(blocksInPrimaryPool, kVSWA_ATTENTION_WINDOW)); + + // Actual release: verify no leaks. + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(0, std::nullopt))); + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); +} + +// Verify multiple consecutive OOW placeholders are handled correctly by storeNewBlock. +// +// With window=4 and tpb=4, the OOW condition (numTokens - removed*4 >= 8) fires twice +// in a single addToken call when numTokens reaches 12 after 11 context tokens: +// 1st OOW: 12 - 0*4 = 12 >= 8 → B0 OOW → P0 inserted +// 2nd OOW: 12 - 1*4 = 8 >= 8 → B1 OOW → P1 inserted +// seq: [P0, P1, B2] +// storeNewBlock fires on the *next* addToken (numTokens=13, usableSize=12): +// blockKeys.size()=3, beam0Blocks=[P0, P1, B2, B3] +// lastBlock=B2(idx 2), prevBlock=P1(idx 1) +// prevBlock->isPlaceholder()==true → "store all blocks" path +// storeBlocks([K0,K1,K2], [P0,P1,B2,B3]): +// insertNodes([K0,K1,K2]) finds/creates all nodes. +// P0 → node K0 has value B0 (storeContextBlocks) → advance prevBlock +// P1 → node K1 has value B1 (storeContextBlocks) → advance prevBlock +// B2 → node K2 is empty → insert +// A subsequent sequence with B0's prefix must reuse B0 (contextCurrentPosition==3). +TEST_F(KVCacheManagerTest, VSWAStoreNewBlockWithMultipleOOWPlaceholders) +{ + auto constexpr blocksInPrimaryPool = 8; + auto const stream = std::make_shared(); + tr::SamplingConfig const samplingConfig{kVSWA_BEAM_WIDTH}; + // window=4 == tpb=4: two OOW events occur before the next storeNewBlock boundary. + auto kvCacheManager = makeSmallWindowManager(blocksInPrimaryPool, stream); + auto const& blockManager = kvCacheManager->getBlockManager(); + SizeType32 constexpr kSmallWindow = 4; + SizeType32 constexpr kSmallTpb = 4; + + // Seq 0: 11 input tokens → B0=[1000..1003], B1=[1004..1007], B2=[1008..1010] (partial). + // storeContextBlocks stores B0 and B1 (both full) so they are in the trie before OOW. + auto inputTokens0 = std::make_shared(11); + std::iota(inputTokens0->begin(), inputTokens0->end(), kVSWA_FIRST_TOKEN); + auto llmRequest0 + = std::make_shared(0, kVSWA_MAX_NEW_TOKENS, inputTokens0, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(0, 11, kVSWA_BEAM_WIDTH, llmRequest0); + kvCacheManager->storeContextBlocks(*llmRequest0); + + // addToken step 1 (token 1011): numTokens=12. + // 12 - 0*4 = 12 >= 4+4=8 → B0 OOW (P0 inserted, numFront=1) + // 12 - 1*4 = 8 >= 8 → B1 OOW (P1 inserted, numFront=2) + // 12 - 2*4 = 4 < 8 → stop + // (12-1)%4=3 != 0 → no new block. seq: [P0, P1, B2] + llmRequest0->addNewToken(kVSWA_FIRST_TOKEN + 11, kVSWA_BEAM_IDX); + kvCacheManager->addToken(0); + kvCacheManager->storeNewBlock(*llmRequest0); // usableSize=11, 11%4!=0 → no-op + + // addToken step 2 (token 1012): numTokens=13. + // 13 - 2*4 = 5 < 8 → no OOW. + // (13-1)%4=0 → B3 allocated. seq: [P0, P1, B2, B3] + // storeNewBlock(usableSize=12): 12%4=0 → fires. + // blockKeys=[K0,K1,K2], beam0Blocks=[P0,P1,B2,B3] + // prevBlock = P1 (index 1) → isPlaceholder()==true → "store all blocks" path + // storeBlocks: P0→advance(B0 in trie); P1→advance(B1 in trie); B2→insert. + llmRequest0->addNewToken(kVSWA_FIRST_TOKEN + 12, kVSWA_BEAM_IDX); + kvCacheManager->addToken(0); + kvCacheManager->storeNewBlock(*llmRequest0); // stores B2 + + // Seq 1: same 4-token prefix as B0 → must reuse B0. + // usableSize for seq1 context = kSmallTpb-1=3 tokens → partial match of 3 tokens. + auto inputTokens1 = std::make_shared(kSmallTpb); + std::iota(inputTokens1->begin(), inputTokens1->end(), kVSWA_FIRST_TOKEN); + auto llmRequest1 + = std::make_shared(1, kVSWA_MAX_NEW_TOKENS, inputTokens1, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(1, kSmallTpb, kVSWA_BEAM_WIDTH, llmRequest1); + EXPECT_EQ(llmRequest1->getContextCurrentPosition(), kSmallTpb - 1); + + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(1, llmRequest1))); + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(0, std::nullopt))); + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); +} + +// Verify that storeBlocks continues past an already-occupied non-placeholder trie slot +// and stores subsequent blocks, leaving the trie intact for later sequences to reuse. +// +// Scenario: +// Seq 0: tokens [K0, K1_A] → B0 at K0, B1_A at K1_A stored in trie on release. +// Seq 1: tokens [K0, K1_B] → reuses B0, allocates B1_B fresh; B1_B stored at K1_B. +// Seq 2: tokens [K0, K1_A] → reuses B0 and B1_A from trie (contextCurrentPosition > 0). +// Seq 2 release storeBlocks sees K0 and K1_A already occupied — skips both cleanly. +// Seq 3: same tokens as seq 2 → must still reuse B0 and B1_A (trie not corrupted). +TEST_F(KVCacheManagerTest, VSWAStoreBlocksSkipsOccupiedSlotsAndContinues) +{ + auto constexpr blocksInPrimaryPool = 8; + auto const stream = std::make_shared(); + tr::SamplingConfig const samplingConfig{kVSWA_BEAM_WIDTH}; + auto kvCacheManager = makeVSWAManager(blocksInPrimaryPool, /*enableBlockReuse=*/true, stream); + auto const& blockManager = kvCacheManager->getBlockManager(); + + TokenIdType constexpr kSharedFirst = kVSWA_FIRST_TOKEN; // shared B0 prefix + TokenIdType constexpr kSeqASecond = kVSWA_FIRST_TOKEN + 100; // B1_A suffix + TokenIdType constexpr kSeqBSecond = kVSWA_FIRST_TOKEN + 200; // B1_B suffix + + // Seq 0: 9 tokens → B0=[1000..1003], B1_A=[1100..1103], partial B2. + auto tokens0 = std::make_shared(9); + std::iota(tokens0->begin(), tokens0->begin() + kVSWA_TOKENS_PER_BLOCK, kSharedFirst); + std::iota(tokens0->begin() + kVSWA_TOKENS_PER_BLOCK, tokens0->end(), kSeqASecond); + auto req0 = std::make_shared(0, kVSWA_MAX_NEW_TOKENS, tokens0, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(0, 9, kVSWA_BEAM_WIDTH, req0); + kvCacheManager->storeContextBlocks(*req0); + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(0, req0))); + // Trie: B0 at K0, B1_A at K1_A under B0. + + // Seq 1: 9 tokens → reuses B0, allocates B1_B for the distinct K1_B suffix. + auto tokens1 = std::make_shared(9); + std::iota(tokens1->begin(), tokens1->begin() + kVSWA_TOKENS_PER_BLOCK, kSharedFirst); + std::iota(tokens1->begin() + kVSWA_TOKENS_PER_BLOCK, tokens1->end(), kSeqBSecond); + auto req1 = std::make_shared(1, kVSWA_MAX_NEW_TOKENS, tokens1, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(1, 9, kVSWA_BEAM_WIDTH, req1); + kvCacheManager->storeContextBlocks(*req1); + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(1, req1))); + // Trie: B0 at K0, B1_A at K1_A and B1_B at K1_B (both children of B0). + + // Seq 2: same prefix as seq 0 ([K0, K1_A, ...]) → reuses B0 and B1_A. + auto tokens2 = std::make_shared(*tokens0); + auto req2 = std::make_shared(2, kVSWA_MAX_NEW_TOKENS, tokens2, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(2, 9, kVSWA_BEAM_WIDTH, req2); + EXPECT_GT(req2->getContextCurrentPosition(), 0); + + // storeBlocks for seq 2: K0 and K1_A are both occupied → skips both without crash. + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(2, req2))); + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); + + // Seq 3: same tokens as seq 2 → B0 and B1_A must still be reusable (trie intact). + auto tokens3 = std::make_shared(*tokens0); + auto req3 = std::make_shared(3, kVSWA_MAX_NEW_TOKENS, tokens3, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(3, 9, kVSWA_BEAM_WIDTH, req3); + EXPECT_GT(req3->getContextCurrentPosition(), 0); + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(3, req3))); + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); +} + +// Verify that storeBlocksForReuse with pinBlocks=true pins both already-in-trie blocks +// (occupied slots, advanced via prevBlock) and newly stored blocks (empty slots). +// After removeSequence the pinned blocks must NOT be in the free pool, and unpinning +// them restores the full pool. +// +// Setup: +// Seq 0: 11 tokens → storeContextBlocks stores B0=[1000..1003] and B1=[1004..1007] +// (both full). B2=[1008..1010] is partial and NOT stored. +// removeSequence with std::nullopt: no storeBlocks; B0 and B1 remain in the +// trie (cached), B2 is freed. +// Seq 1: same 11 tokens → reuses B0 and B1; allocates fresh B2'. +// storeBlocksForReuse(pinBlocks=true): +// usableSize = 10 → blockKeys = [K0_full, K1_full, K2_partial] +// K0 → occupied by B0 → pin B0 (occupied-slot path) +// K1 → occupied by B1 → pin B1 (occupied-slot path) +// K2 → empty → store B2' + pin B2' (empty-slot path) +// pinnedIds.size() == 3. +TEST_F(KVCacheManagerTest, VSWAStoreBlocksForReuseWithPinBlocksPinsAllChainBlocks) +{ + auto constexpr blocksInPrimaryPool = 6; + auto const stream = std::make_shared(); + tr::SamplingConfig const samplingConfig{kVSWA_BEAM_WIDTH}; + auto kvCacheManager = makeVSWAManager(blocksInPrimaryPool, /*enableBlockReuse=*/true, stream); + auto const& blockManager = kvCacheManager->getBlockManager(); + + // Seq 0: 11 tokens → B0 (full), B1 (full) stored by storeContextBlocks; B2 partial. + // Release with std::nullopt so storeBlocks is NOT called → B2 slot stays empty in trie. + auto inputTokens0 = std::make_shared(11); + std::iota(inputTokens0->begin(), inputTokens0->end(), kVSWA_FIRST_TOKEN); + auto llmRequest0 + = std::make_shared(0, kVSWA_MAX_NEW_TOKENS, inputTokens0, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(0, 11, kVSWA_BEAM_WIDTH, llmRequest0); + kvCacheManager->storeContextBlocks(*llmRequest0); + // Release without storing: B0 and B1 remain in the trie; B2 freed. + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(0, std::nullopt))); + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); + + // Seq 1: same 11 tokens → reuses B0 and B1 (contextCurrentPosition > 0); allocates B2'. + auto inputTokens1 = std::make_shared(11); + std::iota(inputTokens1->begin(), inputTokens1->end(), kVSWA_FIRST_TOKEN); + auto llmRequest1 + = std::make_shared(1, kVSWA_MAX_NEW_TOKENS, inputTokens1, samplingConfig, kVSWA_IS_STREAMING); + kvCacheManager->addSequence(1, 11, kVSWA_BEAM_WIDTH, llmRequest1); + EXPECT_GT(llmRequest1->getContextCurrentPosition(), 0); + + // storeBlocksForReuse with pinBlocks=true: + // usableSize=10 → blockKeys=[K0_full, K1_full, K2_partial], beam0Blocks=[B0,B1,B2'] + // K0 occupied by B0 → skip + pin B0. K1 occupied by B1 → skip + pin B1. + // K2 empty → store B2' + pin B2'. Total: 3 pinned blocks. + auto pinnedIds = kvCacheManager->storeBlocksForReuse(1, llmRequest1, /*pinBlocks=*/true); + EXPECT_EQ(static_cast(pinnedIds.size()), 3); + + // removeSequence releases the sequence's ref; pinned blocks keep their extra ref. + EXPECT_NO_THROW(static_cast(kvCacheManager->removeSequence(1, std::nullopt))); + EXPECT_LT(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); + + // Unpinning restores the full pool. + kvCacheManager->unpinBlocksById(pinnedIds); + EXPECT_EQ(blockManager.getNumFreeBlocks(), blocksInPrimaryPool); } diff --git a/cpp/tests/unit_tests/batch_manager/radixBlockTreeTest.cpp b/cpp/tests/unit_tests/batch_manager/radixBlockTreeTest.cpp index def3fa7346be..6838ea906c08 100644 --- a/cpp/tests/unit_tests/batch_manager/radixBlockTreeTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/radixBlockTreeTest.cpp @@ -534,10 +534,10 @@ TEST(MambaTest, kRecurrentStatesSentinelIsNegative) TEST(MambaTest, CreatePlaceholderIsPlaceholder) { - auto ph = KVCacheBlock::createPlaceholder(42); + auto ph = KVCacheBlock::createPlaceholder(); ASSERT_NE(ph, nullptr); EXPECT_TRUE(ph->isPlaceholder()); - EXPECT_EQ(ph->getBlockId(), 42); + EXPECT_EQ(ph->getBlockId(), KVCacheBlock::kPlaceholderBlockId); } TEST(MambaTest, RegularBlockIsNotPlaceholder)