fix kv cache rewind issue

Signed-off-by: Yue Weng <[email protected]>

Author: yweng0828

tensorrt_llm/_torch/pyexecutor/resource_manager.py

@@ -535,14 +535,15 @@ def update_resources(self,
                          scheduled_batch: ScheduledRequests,
                          attn_metadata: "AttentionMetadata" = None,
                          kv_cache_dtype_byte_size: float = None):
-        self.update_kv_cache_draft_token_location(scheduled_batch,
-                                                  attn_metadata,
-                                                  kv_cache_dtype_byte_size)
-        # rewind kv cache
-        for request in scheduled_batch.generation_requests:
-            if request.state != LlmRequestState.GENERATION_COMPLETE:
-                if request.py_rewind_len > 0:
-                    self.rewind_kv_cache(request, request.py_rewind_len)
+        if not self.is_draft:
+            self.update_kv_cache_draft_token_location(scheduled_batch,
+                                                      attn_metadata,
+                                                      kv_cache_dtype_byte_size)
+            # rewind kv cache
+            for request in scheduled_batch.generation_requests:
+                if request.state != LlmRequestState.GENERATION_COMPLETE:
+                    if request.py_rewind_len > 0:
+                        self.rewind_kv_cache(request, request.py_rewind_len)
 
         # For context requests, we store the blocks for reuse.
         for request in scheduled_batch.context_requests:

tensorrt_llm/_torch/pyexecutor/sampler.py

@@ -879,9 +879,9 @@ def _process_draft_tokens_tree(
 
         assert num_accepted_draft_tokens <= longest_accepted_len
 
-        request.py_num_accepted_draft_tokens_indices = eagle_paths[longest_match_path_idx][
-            1:num_accepted_draft_tokens
-        ].tolist()  # exclude the root node
+        tree_node_indices = eagle_paths[longest_match_path_idx][1:num_accepted_draft_tokens]
+        request.py_num_accepted_draft_tokens_indices = (tree_node_indices - 1).tolist()
+
         return num_accepted_draft_tokens - 1
 
     @torch.inference_mode()

tensorrt_llm/_torch/speculative/eagle3.py

@@ -202,7 +202,11 @@ def prepare(self):
             elif is_first_draft and spec_tree_manager is not None:
                 assert req_id in self.request_accepted_path.keys(
                 ), f"Request {req_id} not found in request_accepted_path"
-                accepted_path = self.request_accepted_path[req_id]
+                # 'node_idx + 1' is because we '-1' in sampler.py for kv cache rewind. Now we add it back.
+                accepted_path = [
+                    node_idx + 1
+                    for node_idx in self.request_accepted_path[req_id]
+                ]
 
                 if accepted_path == []:
                     # Case 1: This is a context request, We need to read all the hidden states.
@@ -218,7 +222,6 @@ def prepare(self):
                     assert len(accepted_path_pad) == seq_len
                     hidden_states_read_indices.extend([
                         start_idx + accepted_draft_token_offset
-                        # for accepted_draft_token_offset in accepted_path
                         for accepted_draft_token_offset in accepted_path_pad
                     ])