From 526b3ca235e6c4a3498fc6eeade897882e181750 Mon Sep 17 00:00:00 2001
From: rongwei <scutizhang@tencent.com>
Date: Mon, 17 Nov 2025 13:47:11 +0000
Subject: [PATCH] fix mtp rewind corner case

---
 cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
index 38fde72225..8420c5bb22 100644
--- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
+++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp
@@ -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);