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

@@ -876,9 +876,12 @@ 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
+        # 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()