kanshan/TensorRT-LLMs
1
0
Fork 0
mirror of https://github.com/NVIDIA/TensorRT-LLM.git synced 2026-01-13 22:18:36 +08:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

fix mtp rewind corner case

Browse Source
This commit is contained in:
rongwei 2025-11-17 13:47:11 +00:00 committed by Shi Xiaowei
parent 09d9878385
commit 526b3ca235
1 changed files with 12 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
View File

@ -2861,6 +2861,18 @@ void KVCacheManager::removeToken(RequestIdType requestId)
void KVCacheManager::rewindKVCache(RequestIdType requestId, SizeType32 rewindLengths)
{
// 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
{
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;
}
}
for (SizeType32 si = 0; si < rewindLengths; ++si)
{
removeToken(requestId);