[TRTLLM-4517] [feat] Additional model outputs (#7206)

Signed-off-by: Robin Kobus <19427718+Funatiq@users.noreply.github.com>

Author: Funatiq

examples/llm-api/quickstart_advanced.py

@@ -157,6 +157,12 @@ def add_llm_args(parser):
                         default=False,
                         action='store_true')
     parser.add_argument('--logprobs', default=False, action='store_true')
+
+    parser.add_argument('--additional_model_outputs',
+                        type=str,
+                        default=None,
+                        nargs='+')
+
     return parser
 
 
@@ -279,7 +285,8 @@ def setup_llm(args, **kwargs):
         logprobs=args.logprobs,
         n=args.n,
         best_of=best_of,
-        use_beam_search=use_beam_search)
+        use_beam_search=use_beam_search,
+        additional_model_outputs=args.additional_model_outputs)
     return llm, sampling_params
 
 
@@ -319,6 +326,16 @@ def main():
             if args.logprobs:
                 print(f"[{i}]{sequence_id_text} Logprobs: {sequence.logprobs}")
 
+            if args.additional_model_outputs:
+                for output_name in args.additional_model_outputs:
+                    if sequence.additional_context_outputs:
+                        print(
+                            f"[{i}]{sequence_id_text} Context {output_name}: {sequence.additional_context_outputs[output_name]}"
+                        )
+                    print(
+                        f"[{i}]{sequence_id_text} Generation {output_name}: {sequence.additional_generation_outputs[output_name]}"
+                    )
+
 
 if __name__ == '__main__':
     main()

tensorrt_llm/_torch/pyexecutor/handle_additional_outputs.py

@@ -0,0 +1,112 @@
+from itertools import chain
+from typing import Dict, List
+
+import torch
+
+from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest
+from tensorrt_llm._utils import nvtx_range
+from tensorrt_llm.logger import logger
+
+
+class HandleAdditionalOutputs:
+
+    @torch.inference_mode()
+    @nvtx_range("handle_additional_outputs")
+    def __call__(
+        self,
+        context_requests: List[LlmRequest],
+        generation_requests: List[LlmRequest],
+        outputs: Dict[str, torch.Tensor],
+        beam_width: int,
+        num_context_tokens: int,
+    ):
+        """Handles context and generation logits for a batch of requests.
+
+        Args:
+            context_requests: List of context requests to process
+            generation_requests: List of generation requests to process
+            outputs: Additional outputs tensors
+            beam_width: Beam width for the generation requests
+            num_context_tokens: Number of context tokens in the batch
+        """
+
+        additional_outputs = set()
+        for r in chain(context_requests, generation_requests):
+            if r.py_additional_outputs is not None:
+                additional_outputs.update(r.py_additional_outputs)
+
+        if not additional_outputs:
+            return
+
+        output_length_with_context = num_context_tokens + beam_width * len(
+            generation_requests)
+        output_length_without_context = len(
+            context_requests) + beam_width * len(generation_requests)
+
+        gather_context = {}
+        for name in additional_outputs:
+            if outputs[name].shape[0] == output_length_with_context:
+                gather_context[name] = True
+            else:
+                gather_context[name] = False
+
+        output_index_with_context = 0
+        output_index_without_context = 0
+
+        # Copy additional outputs into decoderBuffers.additional_outputs
+        for llm_req in context_requests:
+            context_output_length = llm_req.context_chunk_size
+
+            outputs_begin = output_index_with_context
+            outputs_end = output_index_with_context + context_output_length
+
+            additional_outputs = llm_req.py_additional_outputs
+            req_context_output = False
+            for name in additional_outputs:
+                if gather_context[name]:
+                    output_device_view = outputs[name][
+                        outputs_begin:outputs_end]
+                    llm_req.py_result.append_additional_context_outputs(
+                        name, output_device_view)
+                    req_context_output = True
+
+            if req_context_output and llm_req.prepopulated_prompt_len > 0:
+                logger.warning(
+                    f"Because of KV cache reuse, not all additional context outputs could be produced for request {llm_req.request_id}."
+                )
+
+            output_index_with_context += context_output_length
+            output_index_without_context += 1
+
+            if llm_req.is_last_context_chunk:
+                for name in additional_outputs:
+                    outputs_begin = (output_index_with_context
+                                     if gather_context[name] else
+                                     output_index_without_context) - 1
+                    outputs_end = outputs_begin + 1
+
+                    output_device_view = outputs[name][
+                        outputs_begin:outputs_end]
+                    llm_req.py_result.append_additional_generation_outputs(
+                        name, torch.tile(output_device_view,
+                                         (1, beam_width, 1)))
+
+        for llm_req in generation_requests:
+            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)
+
+            output_index_with_context += beam_width
+            output_index_without_context += beam_width
+
+        assert output_index_with_context == output_length_with_context, f"output_index_with_context: {output_index_with_context}, output_length_with_context: {output_length_with_context}"
+        assert output_index_without_context == output_length_without_context, f"output_index_without_context: {output_index_without_context}, output_length_without_context: {output_length_without_context}"

tensorrt_llm/_torch/pyexecutor/llm_request.py

@@ -234,7 +234,8 @@ def __init__(self,
                  return_generation_logits: bool = False,
                  exclude_last_generation_logits: bool = False,
                  use_chunked_generation_logits: bool = True,
-                 chunk_size: int = 8):
+                 chunk_size: int = 8,
+                 additional_outputs: Optional[List[str]] = None):
         if streaming and use_chunked_generation_logits:
             assert chunk_size == 1, "chunk_size must be 1 in streaming mode"
         self._streaming = streaming
@@ -253,6 +254,14 @@ def __init__(self,
             chunk_size=self._chunk_size) if return_generation_logits else None
         self._log_probs = LogProbStorage() if return_log_probs else None
         self._mm_embeddings = None
+        self._additional_context_outputs = {
+            name: []
+            for name in additional_outputs
+        } if additional_outputs else None
+        self._additional_generation_outputs = {
+            name: []
+            for name in additional_outputs
+        } if additional_outputs else None
 
     def append_context_logits(self, context_logits: torch.Tensor):
         if self._context_logits:
@@ -277,6 +286,16 @@ def transfer_remaining_device_logits(self):
         if self._generation_logits:
             self._generation_logits.finalize_chunked_transfer()
 
+    def append_additional_context_outputs(
+            self, name: str, additional_context_outputs: torch.Tensor):
+        self._additional_context_outputs[name].append(
+            additional_context_outputs.to("cpu", non_blocking=True))
+
+    def append_additional_generation_outputs(
+            self, name: str, additional_generation_outputs: torch.Tensor):
+        self._additional_generation_outputs[name].append(
+            additional_generation_outputs.to("cpu", non_blocking=True))
+
     def set_log_probs(self, log_probs: list[TokenLogprobs],
                       cum_log_probs: list[float]):
         """
@@ -318,12 +337,37 @@ def cum_log_probs(self) -> list[float] | None:
     def mm_embedding_handle(self) -> Dict[str, Any] | None:
         return self._mm_embeddings
 
+    @property
+    def additional_context_outputs(self) -> Dict[str, torch.Tensor] | None:
+        if self._additional_context_outputs is None:
+            return None
+        outputs = {}
+        for name, output_list in self._additional_context_outputs.items():
+            if len(output_list) == 0:
+                continue
+            outputs[name] = torch.cat(
+                output_list, dim=0) if len(output_list) > 1 else output_list[0]
+        return outputs
+
+    @property
+    def additional_generation_outputs(self) -> Dict[str, torch.Tensor] | None:
+        if self._additional_generation_outputs is None:
+            return None
+        outputs = {}
+        for name, output_list in self._additional_generation_outputs.items():
+            if len(output_list) == 0:
+                continue
+            outputs[name] = torch.cat(
+                output_list, dim=0) if len(output_list) > 1 else output_list[0]
+        return outputs
+
 
 class LlmResult:
     """LlmResult wraps `bindings.executor.Result` but detour some features to Python implementation"""
     py_result_properties = frozenset(
         ('context_logits', 'generation_logits', 'log_probs', 'cum_log_probs',
-         'mm_embedding_handle'))
+         'mm_embedding_handle', 'additional_context_outputs',
+         'additional_generation_outputs'))
 
     def __init__(self,
                  result: Union[bytes, tensorrt_llm.bindings.executor.Result],
@@ -388,6 +432,7 @@ def __init__(
             return_generation_logits: bool = False,
             return_logits_device_memory: bool = True,
             exclude_last_generation_logits: bool = False,
+            additional_outputs: Optional[List[str]] = None,
             return_perf_metrics: bool = False,
             stop_words_list: list[list[int]] | None = None,
             llm_request: Optional[
@@ -448,6 +493,8 @@ def __init__(
         self.py_return_context_logits = return_context_logits
         self.py_return_generation_logits = return_generation_logits
         self.py_return_logits_device_memory = return_logits_device_memory
+        self.py_additional_outputs = additional_outputs
+
         self.py_is_draft = is_draft
         # The request's sequence slot ID, an index between 0 (inclusive) and max_batch_size (exclusive).
         self.py_seq_slot = seq_slot
@@ -477,7 +524,8 @@ def __init__(
             return_generation_logits,
             exclude_last_generation_logits,
             use_chunked_generation_logits=self.py_use_chunked_generation_logits,
-            chunk_size=self.py_logits_chunk_size)
+            chunk_size=self.py_logits_chunk_size,
+            additional_outputs=additional_outputs)
         self.child_requests = []
 
         self._py_embedding_bias_1d: Optional[torch.Tensor] = None
@@ -675,6 +723,11 @@ def executor_request_to_llm_request(
         return_generation_logits=executor_request.output_config.
         return_generation_logits,
         exclude_last_generation_logits=exclude_last_generation_logits,
+        additional_outputs=[
+            output.name for output in
+            executor_request.output_config.additional_model_outputs
+        ] if executor_request.output_config.additional_model_outputs is not None
+        else None,
         draft_tokens=getattr(executor_request, "draft_tokens", None),
         draft_logits=None,
         exclude_input_from_output=executor_request.output_config.

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -2264,21 +2264,34 @@ def _forward_step(self,
         inputs = self._preprocess_inputs(inputs)
         if inputs.get('spec_metadata', None):
             gather_ids = inputs['spec_metadata'].gather_ids
-        if self.without_logits:
-            outputs = self.model_forward(**inputs)
-            return outputs
 
         # For simplicity, just return all the the logits if we have special gather_ids
         # from speculative decoding.
-        logits = self.model_forward(
+        outputs = self.model_forward(
             **inputs,
             return_context_logits=gather_ids is not None
             or gather_context_logits,
         )
-        if gather_ids is not None:
-            return {'logits': logits[gather_ids]}
+
+        if self.without_logits:
+            return outputs
+
+        if isinstance(outputs, dict):
+            # If the model returns a dict, get the logits from it. All other keys are kept.
+            logits = outputs.get('logits', None)
+            # If the logits are not found, no further processing is needed.
+            if logits is None:
+                return outputs
         else:
-            return {'logits': logits}
+            # If the model returns a single tensor, assume it is the logits and wrap it in a dict.
+            logits = outputs
+            outputs = {'logits': logits}
+
+        # If we have special gather_ids, gather the logits
+        if gather_ids is not None:
+            outputs['logits'] = logits[gather_ids]
+
+        return outputs
 
     @nvtx_range("_forward_step_mm_encoder_only")
     def _forward_step_mm_encoder_only(

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -42,6 +42,7 @@
 from ..speculative.drafter import Drafter
 from .executor_request_queue import ExecutorRequestQueue, RequestQueueItem
 from .guided_decoder import GuidedDecoder
+from .handle_additional_outputs import HandleAdditionalOutputs
 from .handle_logits import HandleLogits
 from .kv_cache_connector import KvCacheConnectorManager
 from .kv_cache_transceiver import KvCacheTransceiver
@@ -1815,18 +1816,27 @@ def _sample_async(self, scheduled_batch,
             if batch_outputs is not None:
                 num_context_logits_prefix_sum = [0]
                 prefix_sum = 0
+                num_context_tokens = 0
                 for request in scheduled_batch.context_requests:
-                    prefix_sum += request.context_chunk_size if request.py_return_context_logits else 1
+                    context_chunk_size = request.context_chunk_size
+                    prefix_sum += context_chunk_size if request.py_return_context_logits else 1
                     num_context_logits_prefix_sum.append(prefix_sum)
+                    num_context_tokens += context_chunk_size
+
+                beam_width = self.sampler.beam_width(
+                    scheduled_batch.all_requests())
 
                 HandleLogits()(scheduled_batch.context_requests,
                                scheduled_batch.generation_requests,
-                               batch_outputs["logits"],
-                               self.sampler.beam_width(
-                                   scheduled_batch.all_requests()),
+                               batch_outputs["logits"], beam_width,
                                num_context_logits_prefix_sum,
                                self.sampler.is_generation_model())
 
+                HandleAdditionalOutputs()(scheduled_batch.context_requests,
+                                          scheduled_batch.generation_requests,
+                                          batch_outputs, beam_width,
+                                          num_context_tokens)
+
                 return self.sampler.sample_async(scheduled_batch, batch_outputs,
                                                  num_context_logits_prefix_sum)
         except Exception as e:

tensorrt_llm/executor/result.py

@@ -102,6 +102,8 @@ class CompletionOutput:
         finish_reason (Literal['stop', 'length', 'timeout', 'cancelled'], optional): The reason why the sequence is finished. Defaults to None.
         stop_reason (int, str, optional): The stop string or token id that caused the completion to stop, None if the completion finished for some other reason. Defaults to None.
         generation_logits (torch.Tensor, optional): The logits on the generated output token ids. Defaults to None.
+        additional_context_outputs (Dict[str, torch.Tensor], optional): The additional context outputs. Defaults to None.
+        additional_generation_outputs (Dict[str, torch.Tensor], optional): The additional generation outputs. Defaults to None.
         disaggregated_params (tensorrt_llm.disaggregated_params.DisaggregatedParams, optional): Parameters needed for disaggregated serving. Includes the type of request, the first generated tokens, the context request id and the any additional state needing to be transferred from context and generation instances. Defaults to None.
         request_perf_metrics (tensorrt_llm.bindings.executor.RequestPerfMetrics, optional): Performance metrics for the request. Defaults to None.
 
@@ -122,6 +124,8 @@ class CompletionOutput:
                                     'cancelled']] = None
     stop_reason: Optional[Union[int, str]] = None
     generation_logits: Optional[torch.Tensor] = None
+    additional_context_outputs: Optional[Dict[str, torch.Tensor]] = None
+    additional_generation_outputs: Optional[Dict[str, torch.Tensor]] = None
     disaggregated_params: Optional[DisaggregatedParams] = None
     request_perf_metrics: Optional[tllm.RequestPerfMetrics] = None
 
@@ -387,6 +391,14 @@ def _handle_sequence(self,
             output.generation_logits = response_tensors.generation_logits[
                 src_idx, :output.length]
 
+        if getattr(response_tensors, 'additional_context_outputs',
+                   None) is not None:
+            output.additional_context_outputs = response_tensors.additional_context_outputs
+
+        if getattr(response_tensors, 'additional_generation_outputs',
+                   None) is not None:
+            output.additional_generation_outputs = response_tensors.additional_generation_outputs
+
         # when sampling_params.n > 1 and is cancelled, make sure all the outputs
         # be marked as cancelled.
         if finish_reasons and finish_reasons[