[TRTLLM-10143][feat] Reuse previous draft requests if possible (#10263)

Signed-off-by: ziyixiong-nv <[email protected]>

Author: ziyixiong-nv

tensorrt_llm/_torch/speculative/model_drafter.py

@@ -174,6 +174,16 @@ def _create_accepted_tokens_request(self, request: LlmRequest,
             input_tokens) - num_accepted_tokens - 1
         return new_request
 
+    def _get_previous_draft_request(
+            self, request: LlmRequest) -> Optional[LlmRequest]:
+        """Get the previous draft request for the given request."""
+        if self.previous_draft_batch is None:
+            return None
+        for req in self.previous_draft_batch.all_requests():
+            if req.py_request_id == request.py_request_id:
+                return req
+        return None
+
     def _create_accepted_tokens_request_for_trtllm_attn(
             self, request: LlmRequest, input_tokens: Any,
             num_accepted_tokens: int) -> LlmRequest:
@@ -186,14 +196,24 @@ def _create_accepted_tokens_request_for_trtllm_attn(
         # because at most max_draft_len draft tokens are accepted.
         input_tokens.extend(
             0 for _ in range(self.max_draft_len - num_accepted_tokens))
-        new_request = self._create_draft_request(request, input_tokens)
-        new_request.state = LlmRequestState.GENERATION_IN_PROGRESS
-        new_request.py_num_accepted_draft_tokens = request.py_num_accepted_draft_tokens
-        new_request.py_is_first_draft = True
+
+        # Reuse the previous draft request if it exists.
+        # This can reduce host overhead significantly.
+        draft_request = self._get_previous_draft_request(request)
+        if draft_request is not None:
+            generated_tokens = input_tokens[draft_request.py_prompt_len:]
+            draft_request.set_generated_tokens([generated_tokens])
+        else:
+            draft_request = self._create_draft_request(request, input_tokens)
+
+        draft_request.state = LlmRequestState.GENERATION_IN_PROGRESS
+        draft_request.py_num_accepted_draft_tokens = request.py_num_accepted_draft_tokens
+        draft_request.py_is_first_draft = True
         # For tree decoding, we need to store the accepted tokens indices for these requests,
         # which will be used to update the hidden_states_read_indices.
-        new_request.py_num_accepted_draft_tokens_indices = request.py_num_accepted_draft_tokens_indices
-        return new_request
+        draft_request.py_num_accepted_draft_tokens_indices = request.py_num_accepted_draft_tokens_indices
+
+        return draft_request
 
     def _create_draft_request_for_request(
             self, request: LlmRequest) -> Optional[LlmRequest]: