[fix] Skip logits and additional outputs handling in extra iteration for overlap scheduler

Signed-off-by: Robin Kobus <[email protected]>

Author: Funatiq

tensorrt_llm/_torch/pyexecutor/handle_additional_outputs.py

@@ -3,7 +3,8 @@
 
 import torch
 
-from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest
+from tensorrt_llm._torch.pyexecutor.llm_request import (LlmRequest,
+                                                        LlmRequestState)
 from tensorrt_llm._utils import nvtx_range
 from tensorrt_llm.logger import logger
 
@@ -92,18 +93,19 @@ def __call__(
                                          (1, beam_width, 1)))
 
         for llm_req in generation_requests:
-            additional_outputs = llm_req.py_additional_outputs
+            if llm_req.state != LlmRequestState.GENERATION_COMPLETE:
+                additional_outputs = llm_req.py_additional_outputs
 
-            for name in additional_outputs:
-                outputs_begin = (output_index_with_context
-                                 if gather_context[name] else
-                                 output_index_without_context)
-                outputs_end = outputs_begin + beam_width
-
-                output_device_view = outputs[name][
-                    outputs_begin:outputs_end].reshape(1, beam_width, -1)
-                llm_req.py_result.append_additional_generation_outputs(
-                    name, output_device_view)
+                for name in additional_outputs:
+                    outputs_begin = (output_index_with_context
+                                     if gather_context[name] else
+                                     output_index_without_context)
+                    outputs_end = outputs_begin + beam_width
+
+                    output_device_view = outputs[name][
+                        outputs_begin:outputs_end].reshape(1, beam_width, -1)
+                    llm_req.py_result.append_additional_generation_outputs(
+                        name, output_device_view)
 
             output_index_with_context += beam_width
             output_index_without_context += beam_width

tensorrt_llm/_torch/pyexecutor/handle_logits.py

@@ -3,7 +3,8 @@
 
 import torch
 
-from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest
+from tensorrt_llm._torch.pyexecutor.llm_request import (LlmRequest,
+                                                        LlmRequestState)
 from tensorrt_llm._utils import nvtx_range
 from tensorrt_llm.logger import logger
 
@@ -72,6 +73,9 @@ def __call__(
 
         total_context_logits = num_context_logits_prefix_sum[-1]
         for batch_index, llm_req in enumerate(generation_requests):
+            if llm_req.state == LlmRequestState.GENERATION_COMPLETE:
+                continue
+
             logits_begin = total_context_logits + batch_index * beam_width
             logits_end = logits_begin + beam_width