diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index 64ff4c0d3fc8..a2a367dc3c42 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -40,6 +40,7 @@ #include #include #include +#include #include #include #include @@ -337,8 +338,6 @@ class KVCacheBlock : public std::enable_shared_from_this // Distinct from getPrevBlock() (which navigates the radix lookup tree) BlockPtr mPrevBlockInSeq; - // Iterator pointing to this block in mFreeBlocks. - std::optional mFreeBlockIterator; // Flag indicating if block is full bool mIsFull; @@ -520,6 +519,30 @@ class GenerationRequest std::set mUsedBlocks; // Current prepopulated prompt length SizeType32 mCurrentPrepopulatedPromptLen; + // Serialises per-sequence mutations (addToken, removeToken, storeNewBlock, etc.) + // against removeSequence to prevent use of freed blocks. + mutable std::mutex mMutex; + // Set by removeSequence under mMutex before releasing blocks; checked by all + // mutating callers so they can exit cleanly rather than touch freed state. + bool mRemoved{false}; + +public: + // Acquires mMutex; the returned unique_lock is moveable so callers hold RAII ownership. + [[nodiscard]] std::unique_lock getLock() const + { + return std::unique_lock(mMutex); + } + + bool isRemoved() const noexcept + { + return mRemoved; + } + + // Called by KVCacheManager::removeSequence while holding the lock. + void markRemoved() noexcept + { + mRemoved = true; + } }; // attach metadata to a pool pointer @@ -1608,8 +1631,7 @@ class BaseKVCacheManager // void removeToken(SizeType32 seqSlotIdx); virtual void rewindKVCache(LlmRequest::RequestIdType requestId, SizeType32 rewindLengths) = 0; - [[nodiscard]] virtual GenerationRequest const& getSequence(LlmRequest::RequestIdType requestId) const = 0; - [[nodiscard]] virtual GenerationRequest& getSequence(LlmRequest::RequestIdType requestId) = 0; + [[nodiscard]] virtual std::shared_ptr getSequence(LlmRequest::RequestIdType requestId) const = 0; [[nodiscard]] virtual bool isCrossKv() const = 0; @@ -1982,8 +2004,7 @@ class KVCacheManager : public BaseKVCacheManager void removeToken(LlmRequest::RequestIdType requestId); void rewindKVCache(LlmRequest::RequestIdType requestId, SizeType32 rewindLengths) override; - [[nodiscard]] GenerationRequest const& getSequence(LlmRequest::RequestIdType requestId) const override; - [[nodiscard]] GenerationRequest& getSequence(LlmRequest::RequestIdType requestId) override; + [[nodiscard]] std::shared_ptr getSequence(LlmRequest::RequestIdType requestId) const override; [[nodiscard]] bool isCrossKv() const override { @@ -2104,6 +2125,9 @@ class KVCacheManager : public BaseKVCacheManager SizeType32 sinkTokenLength, SizeType32 blockCapacity, SizeType32 beamWidth, SizeType32 tokensPerBlock); private: + // Removes one token from sequence; caller must hold seqPtr->getLock(). + void removeTokenLocked(GenerationRequest& sequence); + // Maximum number of sequences SizeType32 mMaxNumSequences; // Maximum beam width @@ -2120,7 +2144,7 @@ class KVCacheManager : public BaseKVCacheManager // Block manager BlockManager mBlockManager; // Map of all sequences - std::unordered_map mSequences; + std::unordered_map> mSequences; // Whether to cache KV pages for reuse bool mEnableBlockReuse; // Mutex to protect access to mSequences diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheUtils.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheUtils.h index 90416d124c3a..2c94e556bce4 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheUtils.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheUtils.h @@ -200,7 +200,7 @@ class BlockRange auto windowSize = cacheManager.getBlockManager().getPoolWindowSize(poolIdx); mPoolsPerWindow[windowSize].push_back(cacheManager.getBlockManager().getPrimaryPool(poolIdx)); mBlockIdsPerWindow[windowSize] - = cacheManager.getSequence(mRequestId).getCacheBlockIds(windowSize).at(kFIRST_AND_ONLY_BEAM); + = cacheManager.getSequence(mRequestId)->getCacheBlockIds(windowSize).at(kFIRST_AND_ONLY_BEAM); } if (cacheManager.isEnableIndexerKCache()) { diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 107c703b0b79..c7693fd418a7 100755 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -99,7 +99,6 @@ KVCacheBlock::KVCacheBlock(IdType blockId, tk::KVCacheIndex blockIdx) , mWindowSize{std::numeric_limits::max()} // sentinel: unattached; valid sizes are >= 1 or kRecurrentStates (-1) , mIsPlaceholder{false} - , mFreeBlockIterator(std::nullopt) , mIsFull{false} , mPriority{executor::KvCacheRetentionConfig::kDefaultRetentionPriority} , mDurationMs{std::nullopt} @@ -2273,7 +2272,9 @@ SizeType32 KVCacheManager::getNeededBlocksOneStep( return 0; } - auto const numCurrTokens = getSequence(req.mRequestId).getNumTokens(); + auto const seqPtr_ = getSequence(req.mRequestId); + TLLM_CHECK_WITH_INFO(seqPtr_, "No sequence found for request %lu", req.mRequestId); + auto const numCurrTokens = seqPtr_->getNumTokens(); auto const generatedTokens = numCurrTokens - req.getPromptLen(); auto const maxTokensToAddToKVCache = req.mMaxNewTokens - generatedTokens; auto const tokensPerStep = req.getNumDraftTokens() + 1; @@ -2327,7 +2328,7 @@ SizeType32 KVCacheManager::getRemainingBlocksToCompletion(LlmRequest const& req, if (seqIt != mSequences.end()) { auto const& seq = seqIt->second; - numAllocBlocksPerBeam = seq.getCacheBlockIds(windowSize).at(0).size(); + numAllocBlocksPerBeam = seq->getCacheBlockIds(windowSize).at(0).size(); } } @@ -2440,7 +2441,17 @@ void BlockManager::updateCacheBlockOffsetsAtIdx(GenerationRequest& sequence, Siz void KVCacheManager::addToken(RequestIdType requestId) { // TODO: add streamLLM support - auto& sequence = getSequence(requestId); + auto seqPtr = getSequence(requestId); + if (!seqPtr) + { + return; + } + auto seqLock = seqPtr->getLock(); + if (seqPtr->isRemoved()) + { + return; + } + auto& sequence = *seqPtr; sequence.addNewTokens(1); mBlockManager.adjustBlocksIfNeeded(sequence); } @@ -2502,14 +2513,18 @@ void KVCacheManager::addSequence( ? llmRequest->getKvCacheRetentionConfig().value_or(executor::KvCacheRetentionConfig()) : executor::KvCacheRetentionConfig(); - auto const [seqIt, emplaceDone] = [&] + auto seq = std::make_shared( + requestId, inputLength, beamWidth, mBlockManager.getWindowSizesMetadata(), std::move(kvCacheRetentionConfig)); { auto lck = std::scoped_lock(mSequencesMtx); - return mSequences.try_emplace(requestId, requestId, inputLength, beamWidth, - mBlockManager.getWindowSizesMetadata(), kvCacheRetentionConfig); - }(); - TLLM_CHECK(emplaceDone); - auto& sequence = seqIt->second; + auto const [_, emplaceDone] = mSequences.emplace(requestId, seq); + if (!emplaceDone) + { + return; + } + } + + auto& sequence = *seq; // Get statistics for block allocations/reuse pre request. SizeType32 const numAllocTotalBlocksPreRequest = mBlockManager.getNumAllocTotalBlocks(); @@ -2591,18 +2606,12 @@ void KVCacheManager::addSequence( void KVCacheManager::storeContextBlocks(LlmRequest const& llmRequest) { auto const requestId = llmRequest.mRequestId; - bool found = false; - { - // protect the mSequences - std::scoped_lock lock(mSequencesMtx); - found = mSequences.find(requestId) != mSequences.end(); - } - if (found) + auto seqPtr = getSequence(requestId); + if (seqPtr) { - auto& sequence = getSequence(requestId); if (mEnableBlockReuse && !llmRequest.isDummyRequest()) { - mBlockManager.storeContextBlocks(sequence, llmRequest); + mBlockManager.storeContextBlocks(*seqPtr, llmRequest); } } else @@ -2617,7 +2626,14 @@ void KVCacheManager::storeNewBlock(LlmRequest const& llmRequest) // - Beam search is NOT enabled // - Block reuse is enabled. auto const requestId = llmRequest.mRequestId; - auto& sequence = getSequence(requestId); + auto seqPtr = getSequence(requestId); + TLLM_CHECK_WITH_INFO(seqPtr, "No sequence found for request %lu", requestId); + auto seqLock = seqPtr->getLock(); + if (seqPtr->isRemoved()) + { + return; + } + auto& sequence = *seqPtr; if (sequence.getBeamWidth() > 1 || !mEnableBlockReuse) { return; @@ -2629,21 +2645,30 @@ std::optional KVCacheManager::removeSequence( RequestIdType requestId, OptionalRef llmRequest, bool pinBlocks) { TLLM_LOG_TRACE("[%s]::%s start", isCrossKv() ? "CROSS" : "SELF", __PRETTY_FUNCTION__); - auto sequenceNode = [this, requestId] + auto seqPtr = [this, requestId]() -> std::shared_ptr { std::scoped_lock lock(mSequencesMtx); - return mSequences.extract(requestId); + auto it = mSequences.find(requestId); + if (it == mSequences.end()) + { + return nullptr; + } + auto ptr = it->second; + mSequences.erase(it); + return ptr; }(); std::optional lastStoredId = std::nullopt; - if (!sequenceNode.empty()) + if (seqPtr) { + auto seqLock = seqPtr->getLock(); + seqPtr->markRemoved(); if (mEnableBlockReuse) { - lastStoredId = mBlockManager.releaseBlocks(sequenceNode.mapped(), llmRequest, pinBlocks); + lastStoredId = mBlockManager.releaseBlocks(*seqPtr, llmRequest, pinBlocks); } else { - lastStoredId = mBlockManager.releaseBlocks(sequenceNode.mapped(), std::nullopt, pinBlocks); + lastStoredId = mBlockManager.releaseBlocks(*seqPtr, std::nullopt, pinBlocks); } } if (mBlockManager.isSequenceHeld(requestId)) @@ -2660,7 +2685,14 @@ std::vector KVCacheManager::storeBlocksForReuse( RequestIdType requestId, OptionalRef llmRequest, bool pinBlocks) { TLLM_LOG_TRACE("[%s]::%s start", isCrossKv() ? "CROSS" : "SELF", __PRETTY_FUNCTION__); - auto& sequence = getSequence(requestId); + auto seqPtr = getSequence(requestId); + TLLM_CHECK_WITH_INFO(seqPtr, "No sequence found for request %lu", requestId); + auto seqLock = seqPtr->getLock(); + if (seqPtr->isRemoved()) + { + return {}; + } + auto& sequence = *seqPtr; auto pinnedBlockIds = mBlockManager.storeBlocksForReuse(sequence, llmRequest, pinBlocks); TLLM_LOG_TRACE("[%s]::%s stop", isCrossKv() ? "CROSS" : "SELF", __PRETTY_FUNCTION__); return pinnedBlockIds; @@ -2674,8 +2706,14 @@ void KVCacheManager::schedulingRemoveSequence(RequestIdType requestId) void KVCacheManager::pinBlocks(RequestIdType requestId) { - auto& sequence = getSequence(requestId); - mBlockManager.pinBlocks(sequence); + auto seqPtr = getSequence(requestId); + TLLM_CHECK_WITH_INFO(seqPtr, "No sequence found for request %lu", requestId); + auto seqLock = seqPtr->getLock(); + if (seqPtr->isRemoved()) + { + return; + } + mBlockManager.pinBlocks(*seqPtr); } void KVCacheManager::unpinBlocksById(std::vector const& blockIds) @@ -2705,7 +2743,9 @@ tle::RetentionPriority KVCacheManager::getPriorityByBlockId(KVCacheBlock::IdType SizeType32 KVCacheManager::copyBlockOffsets(ITensor& output, SizeType32 outputSlotOffset, RequestIdType requestId) const { - auto const& sequence = getSequence(requestId); + auto const seqPtr = getSequence(requestId); + TLLM_CHECK_WITH_INFO(seqPtr, "No sequence found for request %lu", requestId); + auto const& sequence = *seqPtr; auto const beamWidth = sequence.getBeamWidth(); auto* dstPtr = bufferCast(output); @@ -3015,10 +3055,8 @@ BlocksPerWindow BaseKVCacheManager::calculateMaxNumBlocks(executor::KvCacheConfi return windowSizeToBlocks; } -void KVCacheManager::removeToken(RequestIdType requestId) +void KVCacheManager::removeTokenLocked(GenerationRequest& sequence) { - // TODO: add streamLLM support - auto& sequence = getSequence(requestId); if (sequence.getNumTokens() == 0) { return; @@ -3035,36 +3073,48 @@ void KVCacheManager::removeToken(RequestIdType requestId) } } -void KVCacheManager::rewindKVCache(RequestIdType requestId, SizeType32 rewindLengths) +void KVCacheManager::removeToken(RequestIdType requestId) { - // Check if the sequence still exists before rewinding - // In overlap mode with MTP, the request may have been terminated and removed - // from mSequences before rewindKVCache is called + // TODO: add streamLLM support + auto seqPtr = getSequence(requestId); + if (!seqPtr) { - std::scoped_lock lck(mSequencesMtx); - if (mSequences.find(requestId) == mSequences.end()) - { - TLLM_LOG_DEBUG("Request %lu has already been removed from KV cache manager, skipping rewind", requestId); - return; - } + return; } - - for (SizeType32 si = 0; si < rewindLengths; ++si) + auto seqLock = seqPtr->getLock(); + if (seqPtr->isRemoved()) { - removeToken(requestId); + return; } + removeTokenLocked(*seqPtr); } -GenerationRequest const& KVCacheManager::getSequence(RequestIdType requestId) const +void KVCacheManager::rewindKVCache(RequestIdType requestId, SizeType32 rewindLengths) { - auto lck = std::scoped_lock(mSequencesMtx); - return mSequences.at(requestId); + // Hold the per-sequence lock for the entire loop so removeSequence cannot + // slip in between iterations and cause a partial rewind. + auto seqPtr = getSequence(requestId); + if (!seqPtr) + { + TLLM_LOG_DEBUG("Request %lu has already been removed from KV cache manager, skipping rewind", requestId); + return; + } + auto seqLock = seqPtr->getLock(); + if (seqPtr->isRemoved()) + { + return; + } + for (SizeType32 si = 0; si < rewindLengths; ++si) + { + removeTokenLocked(*seqPtr); + } } -GenerationRequest& KVCacheManager::getSequence(RequestIdType requestId) +std::shared_ptr KVCacheManager::getSequence(RequestIdType requestId) const { auto lck = std::scoped_lock(mSequencesMtx); - return mSequences.at(requestId); + auto it = mSequences.find(requestId); + return it != mSequences.end() ? it->second : nullptr; } SizeType32 BaseKVCacheManager::getSinkBubbleLength(SizeType32 sinkTokenLen, SizeType32 tokensPerBlock) @@ -3077,7 +3127,9 @@ SizeType32 BaseKVCacheManager::getSinkBubbleLength(SizeType32 sinkTokenLen, Size std::vector> const& KVCacheManager::getCacheBlockIds( RequestIdType requestId, SizeType32 windowSize) const { - return getSequence(requestId).getCacheBlockIds(windowSize); + auto seqPtr = getSequence(requestId); + TLLM_CHECK_WITH_INFO(seqPtr, "No sequence found for request %lu", requestId); + return seqPtr->getCacheBlockIds(windowSize); } std::vector>> KVCacheManager::getBatchCacheBlockIds( @@ -3087,22 +3139,27 @@ std::vector>> KVCacheManager::getBatchCacheB result.reserve(requestIds.size()); for (auto const& requestId : requestIds) { - auto const& sequence = getSequence(requestId); - result.emplace_back(sequence.getCacheBlockIds(windowSize)); + auto const seqPtr = getSequence(requestId); + TLLM_CHECK_WITH_INFO(seqPtr, "No sequence found for request %lu", requestId); + result.emplace_back(seqPtr->getCacheBlockIds(windowSize)); } return result; } std::optional KVCacheManager::getLastBlockId(LlmRequest::RequestIdType requestId) const { - auto const& seq = getSequence(requestId); + auto const seq = getSequence(requestId); + if (!seq) + { + return std::nullopt; + } // Use the first window size auto firstWindowSize = mBlockManager.getFirstWindowSize(); if (firstWindowSize == 0) { return std::nullopt; } - auto const& perBeam = seq.getCacheBlockIds(firstWindowSize); + auto const& perBeam = seq->getCacheBlockIds(firstWindowSize); if (perBeam.empty() || perBeam[0].empty()) { return std::nullopt; diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp index 4c09d8619d9e..6cee09f6d907 100755 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @@ -130,7 +130,7 @@ class PyKvCacheManager : public tbk::BaseKVCacheManager NB_OVERRIDE_PURE(storeBlocksForReuse, requestId, llmRequest, pinBlocks); } - tbk::GenerationRequest const& getSequence(tb::LlmRequest::RequestIdType requestId) const override + std::shared_ptr getSequence(tb::LlmRequest::RequestIdType requestId) const override { NB_OVERRIDE_PURE(getSequence, requestId); } diff --git a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp index 995453fd669d..9b98b4611d3f 100755 --- a/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp +++ b/cpp/tests/unit_tests/batch_manager/kvCacheManagerTest.cpp @@ -657,7 +657,7 @@ TEST_F(KVCacheManagerTest, FindBlocksInReuseTreeByBlockKeysTest) auto llmRequest0 = std::make_shared(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming); kvCacheManager.addSequence(requestId, inputLength, beamWidth, llmRequest0); EXPECT_EQ(llmRequest0->getContextCurrentPosition(), 0); - auto cacheBlockIds = kvCacheManager.getSequence(requestId).getCacheBlockIds(maxAttentionWindow).at(beamIdx); + auto cacheBlockIds = kvCacheManager.getSequence(requestId)->getCacheBlockIds(maxAttentionWindow).at(beamIdx); EXPECT_THAT(cacheBlockIds, ::testing::ElementsAreArray({0, 1, 2})); (void) kvCacheManager.removeSequence(requestId, llmRequest0); @@ -2795,7 +2795,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerLeafBlockTest) auto llmRequest1 = std::make_shared(1, maxNewTokens, inputTokens1, samplingConfig, isStreaming); kvCacheManager.addSequence(1, inputLength1, beamWidth, llmRequest1); - GenerationRequest const& seq1 = kvCacheManager.getSequence(1); + auto& seq1 = *kvCacheManager.getSequence(1); EXPECT_EQ(llmRequest1->getContextCurrentPosition(), 0); // Block 1 should NOT be reused. It was not freed even if partial. EXPECT_THAT(seq1.getCacheBlockIds(maxAttentionWindow).at(0), ::testing::ElementsAreArray({2})); @@ -2866,7 +2866,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerLeafBlockWithDependentTest) = std::make_shared(requestId0, maxNewTokens, inputTokens0, samplingConfig, isStreaming); kvCacheManager.addSequence(requestId0, inputLength0, beamWidth, llmRequest0); kvCacheManager.storeContextBlocks(*llmRequest0); - GenerationRequest const& seq0 = kvCacheManager.getSequence(requestId0); + auto& seq0 = *kvCacheManager.getSequence(requestId0); // Add two more blocks of generated tokens for (int i = tokensPerBlock; i < 3 * tokensPerBlock; ++i) @@ -2893,7 +2893,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerLeafBlockWithDependentTest) = std::make_shared(requestId1, maxNewTokens, inputTokens1, samplingConfig, isStreaming); kvCacheManager.addSequence(requestId1, inputLength1, beamWidth, llmRequest1); kvCacheManager.storeContextBlocks(*llmRequest1); - GenerationRequest const& seq1 = kvCacheManager.getSequence(requestId1); + auto& seq1 = *kvCacheManager.getSequence(requestId1); // Verify that all primary blocks are in use EXPECT_EQ(blockManager.getNumFreeBlocks(), 0); @@ -3418,7 +3418,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerMaxAttentionWindowSmallerThanBlockSizeT auto constexpr beamIdx = 0; kvCacheManager.addSequence(requestId, inputLength, beamWidth, llmRequest); - GenerationRequest const& seq0 = kvCacheManager.getSequence(requestId); + auto& seq0 = *kvCacheManager.getSequence(requestId); EXPECT_EQ(llmRequest->getContextCurrentPosition(), 0); EXPECT_THAT(seq0.getCacheBlockIds(onlyWindowSize).at(beamIdx), ::testing::ElementsAreArray({0})); @@ -3688,7 +3688,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerMaxAttentionWindowWithReuseTest) auto constexpr beamIdx = 0; kvCacheManager.addSequence(requestId, inputLength, beamWidth, llmRequest); - GenerationRequest const& seq0 = kvCacheManager.getSequence(requestId); + auto& seq0 = *kvCacheManager.getSequence(requestId); EXPECT_EQ(llmRequest->getContextCurrentPosition(), 0); EXPECT_THAT(seq0.getCacheBlockIds(onlyWindowSize).at(beamIdx), ::testing::ElementsAreArray({0, 1, 2, 3})); @@ -3719,7 +3719,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerMaxAttentionWindowWithReuseTest) inputTokens->resize(inputLength); llmRequest = std::make_shared(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming); kvCacheManager.addSequence(requestId, inputLength, beamWidth, llmRequest); - GenerationRequest const& seq1 = kvCacheManager.getSequence(requestId); + auto& seq1 = *kvCacheManager.getSequence(requestId); EXPECT_EQ(llmRequest->getContextCurrentPosition(), 6); EXPECT_THAT(seq1.getCacheBlockIds(onlyWindowSize).at(beamIdx), ::testing::ElementsAreArray( @@ -3742,7 +3742,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerMaxAttentionWindowWithReuseTest) std::iota(inputTokens->begin(), inputTokens->end(), firstToken); llmRequest = std::make_shared(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming); kvCacheManager.addSequence(requestId, inputLength, beamWidth, llmRequest); - GenerationRequest const& seq2 = kvCacheManager.getSequence(requestId); + auto& seq2 = *kvCacheManager.getSequence(requestId); EXPECT_EQ(llmRequest->getContextCurrentPosition(), 9); EXPECT_THAT(seq2.getCacheBlockIds(onlyWindowSize).at(beamIdx), ::testing::ElementsAreArray({0, 1, 7})); EXPECT_NO_THROW(static_cast(kvCacheManager.removeSequence(requestId, llmRequest))); @@ -3757,7 +3757,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerMaxAttentionWindowWithReuseTest) std::iota(inputTokens->begin(), inputTokens->end(), firstToken); llmRequest = std::make_shared(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming); kvCacheManager.addSequence(requestId, inputLength, beamWidth, llmRequest); - GenerationRequest const& seq3 = kvCacheManager.getSequence(requestId); + auto& seq3 = *kvCacheManager.getSequence(requestId); EXPECT_EQ(llmRequest->getContextCurrentPosition(), 14); EXPECT_THAT(seq3.getCacheBlockIds(onlyWindowSize).at(beamIdx), ::testing::ElementsAreArray({0, 1, 2, 8})); @@ -3821,7 +3821,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerSWAInvalidateReuseTest) = std::make_shared(/*requestId=*/0, maxNewTokens, inputTokens, samplingConfig, isStreaming); kvCacheManager.addSequence(/*requestId=*/0, inputLength, beamWidth, llmRequest0); - GenerationRequest const& seq0 = kvCacheManager.getSequence(/*requestId=*/0); + auto& seq0 = *kvCacheManager.getSequence(/*requestId=*/0); // Block 0 goes out-of-window and is detached llmRequest0->addNewToken(1011, beamIdx); @@ -3839,7 +3839,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerSWAInvalidateReuseTest) auto llmRequest1 = std::make_shared(/*requestId=*/1, maxNewTokens, inputTokens, samplingConfig, isStreaming); kvCacheManager.addSequence(/*requestId=*/1, inputLength, beamWidth, llmRequest1); - GenerationRequest const& seq1 = kvCacheManager.getSequence(/*requestId=*/1); + auto& seq1 = *kvCacheManager.getSequence(/*requestId=*/1); auto const onlyWindowSize = theOnlyWindowSize(kvCacheManager); EXPECT_FALSE(blockManager.isSequenceValidForStoreForReuse(seq0.getRequestId(), onlyWindowSize)); @@ -3902,7 +3902,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerVariableWindowAttentionWithReuseTest) auto constexpr beamIdx = 0; auto const assertBlocks - = [minAttentionWindow, maxAttentionWindow, beamIdx](GenerationRequest seq, + = [minAttentionWindow, maxAttentionWindow, beamIdx](GenerationRequest const& seq, std::initializer_list expectedBlocksMin, std::initializer_list expectedBlocksMax) { auto blocksMin = seq.getCacheBlockIds(minAttentionWindow).at(beamIdx); @@ -3920,7 +3920,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerVariableWindowAttentionWithReuseTest) auto llmRequest = std::make_shared(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming); kvCacheManager.addSequence(requestId, inputLength, beamWidth, llmRequest); - GenerationRequest const& seq0 = kvCacheManager.getSequence(requestId); + auto& seq0 = *kvCacheManager.getSequence(requestId); EXPECT_EQ(llmRequest->getContextCurrentPosition(), 0); assertBlocks(seq0, {0, 1}, {0, 1}); @@ -3952,7 +3952,7 @@ TEST_F(KVCacheManagerTest, KVCacheManagerVariableWindowAttentionWithReuseTest) std::iota(inputTokens->begin(), inputTokens->end(), firstToken); llmRequest = std::make_shared(requestId, maxNewTokens, inputTokens, samplingConfig, isStreaming); kvCacheManager.addSequence(requestId, inputLength, beamWidth, llmRequest); - GenerationRequest const& seq1 = kvCacheManager.getSequence(requestId); + auto& seq1 = *kvCacheManager.getSequence(requestId); EXPECT_EQ(llmRequest->getContextCurrentPosition(), 6); assertBlocks(seq1, {0, 3}, {0, 3}); @@ -4137,7 +4137,7 @@ TEST_F(KVCacheManagerTest, GetPriorityByBlockId) kvCacheManager.storeContextBlocks(*llmRequest); // Get block IDs for the sequence - auto const& seq = kvCacheManager.getSequence(0); + auto& seq = *kvCacheManager.getSequence(0); auto cacheBlockIds = seq.getCacheBlockIds(maxAttentionWindow).at(0); ASSERT_GE(cacheBlockIds.size(), 1);