From 8846dac9b46ed1f6e86d7d4761c3b11b65be566a Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Wed, 29 Oct 2025 02:42:23 +0800 Subject: [PATCH] [https://nvbugs/5578175][fix] Fix block range index (#8470) Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Signed-off-by: Mike Iovine <6158008+mikeiovine@users.noreply.github.com> Signed-off-by: Mike Iovine --- cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index aaae98ff2a..b69db6d1bc 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -2145,7 +2145,7 @@ SizeType32 KVCacheManager::getNeededBlocksOneStep( return 0; } - auto const numCurrTokens = mSequences.at(req.mRequestId).getNumTokens(); + auto const numCurrTokens = getSequence(req.mRequestId).getNumTokens(); auto const generatedTokens = numCurrTokens - req.getPromptLen(); auto const maxTokensToAddToKVCache = req.mMaxNewTokens - generatedTokens; auto const tokensPerStep = req.getNumDraftTokens() + 1; @@ -2409,7 +2409,13 @@ void KVCacheManager::addSequence( void KVCacheManager::storeContextBlocks(LlmRequest const& llmRequest) { auto const requestId = llmRequest.mRequestId; - if (mSequences.find(requestId) != mSequences.end()) + bool found = false; + { + // protect the mSequences + std::scoped_lock lock(mSequencesMtx); + found = mSequences.find(requestId) != mSequences.end(); + } + if (found) { auto& sequence = getSequence(requestId); if (mEnableBlockReuse && !llmRequest.isDummyRequest())