[TRTLLM-5670][feat] Support top log probabilities

cpp/include/tensorrt_llm/executor/executor.h

@@ -215,7 +215,8 @@ class OutputConfig
     explicit OutputConfig(bool returnLogProbs = false, bool returnContextLogits = false,
         bool returnGenerationLogits = false, bool excludeInputFromOutput = false, bool returnEncoderOutput = false,
         bool returnPerfMetrics = false,
-        std::optional<std::vector<AdditionalModelOutput>> additionalModelOutputs = std::nullopt);
+        std::optional<std::vector<AdditionalModelOutput>> additionalModelOutputs = std::nullopt,
+        std::optional<SizeType32> topLogProbs = std::nullopt);
 
     /// @brief Controls if Result should contain log probabilities. Default is false.
     bool returnLogProbs;
@@ -233,6 +234,8 @@ class OutputConfig
 
     /// @brief The additional outputs to gather from the model.
     std::optional<std::vector<AdditionalModelOutput>> additionalModelOutputs;
+    /// @brief The number of logprobs to return. Default is 0.
+    std::optional<SizeType32> topLogProbs;
 };
 
 /// @brief Configuration for speculative decoding with external draft tokens.

cpp/tensorrt_llm/executor/outputConfig.cpp

@@ -22,14 +22,15 @@ namespace tensorrt_llm::executor
 
 OutputConfig::OutputConfig(bool inReturnLogProbs, bool inReturnContextLogits, bool inReturnGenerationLogits,
     bool inExcludeInputFromOutput, bool inReturnEncoderOutput, bool inReturnPerfMetrics,
-    std::optional<std::vector<AdditionalModelOutput>> additionalModelOutputs)
+    std::optional<std::vector<AdditionalModelOutput>> additionalModelOutputs, std::optional<SizeType32> topLogProbs)
     : returnLogProbs(inReturnLogProbs)
     , returnContextLogits(inReturnContextLogits)
     , returnGenerationLogits(inReturnGenerationLogits)
     , excludeInputFromOutput(inExcludeInputFromOutput)
     , returnEncoderOutput(inReturnEncoderOutput)
     , returnPerfMetrics(inReturnPerfMetrics)
     , additionalModelOutputs(std::move(additionalModelOutputs))
+    , topLogProbs(topLogProbs)
 {
 }
 

cpp/tensorrt_llm/nanobind/executor/request.cpp

@@ -198,42 +198,46 @@ void initRequestBindings(nb::module_& m)
     auto outputConfigGetstate = [](tle::OutputConfig const& self)
     {
         return nb::make_tuple(self.returnLogProbs, self.returnContextLogits, self.returnGenerationLogits,
-            self.excludeInputFromOutput, self.returnEncoderOutput, self.returnPerfMetrics, self.additionalModelOutputs);
+            self.excludeInputFromOutput, self.returnEncoderOutput, self.returnPerfMetrics, self.additionalModelOutputs,
+            self.topLogProbs);
     };
     auto outputConfigSetstate = [](tle::OutputConfig& outputConfig, nb::tuple const& state)
     {
-        if (state.size() != 7)
+        if (state.size() != 8)
         {
             throw std::runtime_error("Invalid OutputConfig state!");
         }
         new (&outputConfig) tle::OutputConfig(nb::cast<bool>(state[0]), nb::cast<bool>(state[1]),
             nb::cast<bool>(state[2]), nb::cast<bool>(state[3]), nb::cast<bool>(state[4]), nb::cast<bool>(state[5]),
-            nb::cast<std::optional<std::vector<tle::AdditionalModelOutput>>>(state[6]));
+            nb::cast<std::optional<std::vector<tle::AdditionalModelOutput>>>(state[6]),
+            nb::cast<std::optional<SizeType32>>(state[7]));
     };
     nb::class_<tle::OutputConfig>(m, "OutputConfig")
         .def(
             "__init__",
             [](tle::OutputConfig& self, std::optional<bool> return_log_probs, std::optional<bool> return_context_logits,
                 std::optional<bool> return_generation_logits, std::optional<bool> exclude_input_from_output,
                 std::optional<bool> return_encoder_output, std::optional<bool> return_perf_metrics,
-                std::optional<std::vector<tle::AdditionalModelOutput>> additional_model_outputs)
+                std::optional<std::vector<tle::AdditionalModelOutput>> additional_model_outputs,
+                std::optional<SizeType32> top_logprobs)
             {
                 new (&self) tle::OutputConfig(return_log_probs.value_or(false), return_context_logits.value_or(false),
                     return_generation_logits.value_or(false), exclude_input_from_output.value_or(false),
                     return_encoder_output.value_or(false), return_perf_metrics.value_or(false),
-                    additional_model_outputs);
+                    additional_model_outputs, top_logprobs);
             },
             nb::arg("return_log_probs") = nb::none(), nb::arg("return_context_logits") = nb::none(),
             nb::arg("return_generation_logits") = nb::none(), nb::arg("exclude_input_from_output") = nb::none(),
             nb::arg("return_encoder_output") = nb::none(), nb::arg("return_perf_metrics") = nb::none(),
-            nb::arg("additional_model_outputs") = nb::none())
+            nb::arg("additional_model_outputs") = nb::none(), nb::arg("top_logprobs") = nb::none())
         .def_rw("return_log_probs", &tle::OutputConfig::returnLogProbs)
         .def_rw("return_context_logits", &tle::OutputConfig::returnContextLogits)
         .def_rw("return_generation_logits", &tle::OutputConfig::returnGenerationLogits)
         .def_rw("exclude_input_from_output", &tle::OutputConfig::excludeInputFromOutput)
         .def_rw("return_encoder_output", &tle::OutputConfig::returnEncoderOutput)
         .def_rw("return_perf_metrics", &tle::OutputConfig::returnPerfMetrics)
         .def_rw("additional_model_outputs", &tle::OutputConfig::additionalModelOutputs)
+        .def_rw("top_logprobs", &tle::OutputConfig::topLogProbs)
         .def("__getstate__", outputConfigGetstate)
         .def("__setstate__", outputConfigSetstate);
 

cpp/tensorrt_llm/pybind/executor/request.cpp

@@ -191,31 +191,35 @@ void initRequestBindings(pybind11::module_& m)
     auto outputConfigGetstate = [](tle::OutputConfig const& self)
     {
         return py::make_tuple(self.returnLogProbs, self.returnContextLogits, self.returnGenerationLogits,
-            self.excludeInputFromOutput, self.returnEncoderOutput, self.returnPerfMetrics, self.additionalModelOutputs);
+            self.excludeInputFromOutput, self.returnEncoderOutput, self.returnPerfMetrics, self.additionalModelOutputs,
+            self.topLogProbs);
     };
     auto outputConfigSetstate = [](py::tuple const& state)
     {
-        if (state.size() != 7)
+        if (state.size() != 8)
         {
             throw std::runtime_error("Invalid OutputConfig state!");
         }
         return tle::OutputConfig(state[0].cast<bool>(), state[1].cast<bool>(), state[2].cast<bool>(),
             state[3].cast<bool>(), state[4].cast<bool>(), state[5].cast<bool>(),
-            state[6].cast<std::optional<std::vector<tle::AdditionalModelOutput>>>());
+            state[6].cast<std::optional<std::vector<tle::AdditionalModelOutput>>>(),
+            state[7].cast<std::optional<SizeType32>>());
     };
     py::class_<tle::OutputConfig>(m, "OutputConfig")
-        .def(py::init<bool, bool, bool, bool, bool, bool, std::optional<std::vector<tle::AdditionalModelOutput>>>(),
+        .def(py::init<bool, bool, bool, bool, bool, bool, std::optional<std::vector<tle::AdditionalModelOutput>>,
+                 std::optional<SizeType32>>(),
             py::arg("return_log_probs") = false, py::arg("return_context_logits") = false,
             py::arg("return_generation_logits") = false, py::arg("exclude_input_from_output") = false,
             py::arg("return_encoder_output") = false, py::arg("return_perf_metrics") = false,
-            py::arg("additional_model_outputs") = py::none())
+            py::arg("additional_model_outputs") = py::none(), py::arg("top_logprobs") = py::none())
         .def_readwrite("return_log_probs", &tle::OutputConfig::returnLogProbs)
         .def_readwrite("return_context_logits", &tle::OutputConfig::returnContextLogits)
         .def_readwrite("return_generation_logits", &tle::OutputConfig::returnGenerationLogits)
         .def_readwrite("exclude_input_from_output", &tle::OutputConfig::excludeInputFromOutput)
         .def_readwrite("return_encoder_output", &tle::OutputConfig::returnEncoderOutput)
         .def_readwrite("return_perf_metrics", &tle::OutputConfig::returnPerfMetrics)
         .def_readwrite("additional_model_outputs", &tle::OutputConfig::additionalModelOutputs)
+        .def_readwrite("top_logprobs", &tle::OutputConfig::topLogProbs)
         .def(py::pickle(outputConfigGetstate, outputConfigSetstate));
 
     auto externalDraftTokensConfigGetstate = [](tle::ExternalDraftTokensConfig const& self)

examples/llm-api/quickstart_advanced.py

@@ -147,6 +147,8 @@ def add_llm_args(parser):
                         default=False,
                         action='store_true')
     parser.add_argument('--logprobs', default=False, action='store_true')
+    parser.add_argument('--top_logprobs', type=int, default=None)
+    parser.add_argument('--max_top_logprobs', type=int, default=0)
     return parser
 
 
@@ -213,6 +215,14 @@ def setup_llm(args, **kwargs):
         batching_wait_iters=args.attention_dp_batching_wait_iters,
     )
 
+    # if only top_logprobs is set, set max_top_logprobs to top_logprobs
+    if args.top_logprobs is not None and args.max_top_logprobs < args.top_logprobs:
+        args.max_top_logprobs = args.top_logprobs
+    # Remove this once torch sampler stops using enable_mixed_sampler
+    is_greedy = (not args.max_beam_width > 1) and (
+        args.top_k is None or args.top_k == 1) and (args.top_p is None
+                                                    or args.top_p == 0.0)
+
     llm = LLM(
         model=args.model_dir,
         backend='pytorch',
@@ -246,6 +256,8 @@ def setup_llm(args, **kwargs):
         trust_remote_code=args.trust_remote_code,
         gather_generation_logits=args.return_generation_logits,
         max_beam_width=args.max_beam_width,
+        max_top_logprobs=args.max_top_logprobs,
+        enable_mixed_sampler=not is_greedy,
         **kwargs)
 
     use_beam_search = args.max_beam_width > 1
@@ -257,14 +269,19 @@ def setup_llm(args, **kwargs):
 
     assert best_of >= args.n, f"In sampling mode best_of value: {best_of} should be less or equal to n: {args.n}"
 
+    num_logprobs = args.logprobs
+    if args.logprobs is not None:
+        if args.top_logprobs is not None:
+            num_logprobs = args.top_logprobs
+
     sampling_params = SamplingParams(
         max_tokens=args.max_tokens,
         temperature=args.temperature,
         top_k=args.top_k,
         top_p=args.top_p,
         return_context_logits=args.return_context_logits,
         return_generation_logits=args.return_generation_logits,
-        logprobs=args.logprobs,
+        logprobs=num_logprobs,
         n=args.n,
         best_of=best_of,
         use_beam_search=use_beam_search)

tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py

@@ -345,6 +345,7 @@ def create_autodeploy_executor(ad_config: LlmArgs):
         max_seq_len=ad_config.max_seq_len,
         max_draft_len=max_draft_len,
         max_num_sequences=max_num_sequences,
+        max_top_logprobs=ad_config.max_top_logprobs,
         max_beam_width=ad_config.max_beam_width,
         enable_mixed_sampler=ad_config.enable_mixed_sampler,
     )

tensorrt_llm/_torch/auto_deploy/shim/demollm.py

@@ -234,9 +234,9 @@ def _sample(
         logits_shape = logits.shape
         logits = logits.view(-1, logits_shape[-1])  # sampling_batch expects 2D logits
         if isinstance(sampling_params.top_k, int):
-            idx_next, probs = top_k_sampling_batch(logits, sampling_params.top_k)
+            idx_next, probs, _ = top_k_sampling_batch(logits, sampling_params.top_k)
         else:
-            idx_next, probs = greedy_search_sampling_batch(logits)
+            idx_next, probs, _ = greedy_search_sampling_batch(logits)
         idx_next = idx_next.view(logits_shape[:-1])
         return idx_next, probs
 

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -686,7 +686,8 @@ def create_py_executor_instance(
 
 
 def create_torch_sampler_args(mapping: Mapping, *, max_seq_len: int,
-                              enable_mixed_sampler: bool, max_batch_size: int,
+                              enable_mixed_sampler: bool, max_top_logprobs: int,
+                              max_batch_size: int,
                               speculative_config: SpeculativeConfig,
                               max_beam_width: int):
     max_num_sequences = max_batch_size * mapping.pp_size
@@ -697,6 +698,7 @@ def create_torch_sampler_args(mapping: Mapping, *, max_seq_len: int,
         max_draft_len=max_draft_len,
         max_num_sequences=max_num_sequences,
         max_beam_width=max_beam_width,
+        max_top_logprobs=max_top_logprobs,
         enable_mixed_sampler=enable_mixed_sampler,
     )
 
@@ -712,6 +714,7 @@ def instantiate_sampler(engine: PyTorchModelEngine,
         mapping,
         max_seq_len=engine.max_seq_len,
         enable_mixed_sampler=pytorch_backend_config.enable_mixed_sampler,
+        max_top_logprobs=pytorch_backend_config.max_top_logprobs,
         max_batch_size=max_batch_size,
         speculative_config=speculative_config,
         max_beam_width=max_beam_width)

tensorrt_llm/_torch/pyexecutor/config.py

@@ -113,6 +113,9 @@ class PyTorchConfig:
     # If false, set the PyTorch CUDA memory fraction to 1.0.
     _limit_torch_cuda_mem_fraction: bool = True
 
+    # The max number of logprobs to return per token
+    max_top_logprobs: int = 0
+
 
 EXETENDED_EXECUTOR_CONFIG_FIELDS = [
     'backend',

tensorrt_llm/_torch/pyexecutor/llm_request.py

@@ -117,16 +117,21 @@ def set_exclude_last(self, should_exclude_last: bool) -> None:
 class LogProbStorage:
     beam_width: int = -1
     log_probs: list[TokenLogprobs]
+    top_log_probs: list[TokenLogprobs]
     cum_log_probs: list[float]
 
     def _init(self, first_input: list[TokenLogprobs]):
         self.beam_width = len(first_input)
         self.log_probs = [[] for _ in range(self.beam_width)]
         self.cum_log_probs = [0 for _ in range(self.beam_width)]
-
-    def append(self,
-               new_probs: list[TokenLogprobs],
-               cum_log_probs: Optional[list[float]] = None):
+        self.top_log_probs = [[] for _ in range(self.beam_width)]
+
+    def append(
+        self,
+        new_probs: list[TokenLogprobs],
+        cum_log_probs: Optional[list[float]] = None,
+        top_log_probs: Optional[list[TokenLogprobs]] = None,
+    ):
         """
         new_probs: [beam_width, num_tokens]
         cum_log_probs: [beam_width]
@@ -137,14 +142,18 @@ def append(self,
         assert len(new_probs) == self.beam_width, "Beam width mismatch"
         for beam_idx, probs in enumerate(new_probs):
             self.log_probs[beam_idx].extend(probs)
+            if top_log_probs is not None:
+                self.top_log_probs[beam_idx].extend(top_log_probs[beam_idx])
             if cum_log_probs is not None:
                 self.cum_log_probs[beam_idx] = cum_log_probs[beam_idx]
             else:
                 self.cum_log_probs[beam_idx] += sum(
                     next(iter(prob.values())).logprob for prob in probs)
 
-    def set_log_probs(self, log_probs: list[TokenLogprobs],
-                      cum_log_probs: list[float]):
+    def set_log_probs(self,
+                      log_probs: list[TokenLogprobs],
+                      cum_log_probs: list[float],
+                      top_log_probs: Optional[list[TokenLogprobs]] = None):
         """
         Reset the storage and refill it with new values
         log_probs: [beam_width, num_tokens]
@@ -153,7 +162,7 @@ def set_log_probs(self, log_probs: list[TokenLogprobs],
         # reinitialize the storage to clear the lists
         self._init(log_probs)
         # append the new values
-        self.append(log_probs, cum_log_probs)
+        self.append(log_probs, cum_log_probs, top_log_probs)
 
 
 class PyResult:
@@ -187,23 +196,34 @@ def append_generation_logits(self, generation_logits: torch.Tensor):
 
     def append_log_probs(self,
                          log_probs: list[TokenLogprobs],
-                         cum_log_probs: Optional[list[float]] = None):
+                         cum_log_probs: Optional[list[float]] = None,
+                         top_log_probs: Optional[list[TokenLogprobs]] = None):
+        """
+        Append log_probs and optionally cum_log_probs and top_log_probs
+        log_probs: [beam_width, num_tokens]
+        cum_log_probs: [beam_width]
+        top_log_probs: [beam_width, num_tokens]
+        """
         if self._log_probs:
-            self._log_probs.append(log_probs, cum_log_probs)
+            self._log_probs.append(log_probs, cum_log_probs, top_log_probs)
 
     def append_mm_embeddings(self, mm_embeddings: torch.Tensor):
         self._mm_embeddings = SharedTensorContainer.from_tensor(
             mm_embeddings).dump_to_dict()
 
-    def set_log_probs(self, log_probs: list[TokenLogprobs],
-                      cum_log_probs: list[float]):
+    def set_log_probs(self,
+                      log_probs: list[TokenLogprobs],
+                      cum_log_probs: list[float],
+                      top_log_probs: Optional[list[TokenLogprobs]] = None):
         """
-        Set log_probs and cum_log_probs to the new values
+        Set log_probs, cum_log_probs and top_log_probs to the new values
         log_probs: [beam_width, num_tokens]
         cum_log_probs: [beam_width]
+        top_log_probs: [beam_width, num_tokens]
         """
         if self._log_probs:
-            self._log_probs.set_log_probs(log_probs, cum_log_probs)
+            self._log_probs.set_log_probs(log_probs, cum_log_probs,
+                                          top_log_probs)
 
     @property
     def context_logits(self) -> torch.Tensor | None:
@@ -228,6 +248,10 @@ def generation_logits(self) -> torch.Tensor | None:
     def log_probs(self) -> list[TokenLogprobs] | None:
         return self._log_probs and self._log_probs.log_probs
 
+    @property
+    def top_log_probs(self) -> list[TokenLogprobs] | None:
+        return self._log_probs and self._log_probs.top_log_probs
+
     @property
     def cum_log_probs(self) -> list[float] | None:
         return self._log_probs and self._log_probs.cum_log_probs
@@ -241,7 +265,7 @@ 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'))
+         'top_log_probs', 'mm_embedding_handle'))
 
     def __init__(self,
                  result: Union[bytes, tensorrt_llm.bindings.executor.Result],
@@ -302,6 +326,7 @@ def __init__(
             is_draft: bool = False,
             seq_slot: Optional[int] = None,
             target_seq_slot: Optional[int] = None,
+            top_logprobs: Optional[int] = None,
             **kwargs):
 
         self.py_logits_post_processors = kwargs.pop("py_logits_post_processors",
@@ -346,6 +371,7 @@ def __init__(
             TaskLayerModuleConfig] | None = None
 
         self.py_return_log_probs = return_log_probs
+        self.py_top_logprobs = top_logprobs
         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
@@ -394,6 +420,10 @@ def create_response(
     def is_dummy(self):
         return self.is_attention_dp_dummy or self.is_cuda_graph_dummy or self.is_dummy_request
 
+    @property
+    def calculate_top_logprobs(self):
+        return self.py_top_logprobs is not None and self.py_top_logprobs > 0
+
     def finish_by(self, reason: FinishReason, beam: int) -> None:
         """CPP finish by reason does not support beam_width > 1"""
         self.state = LlmRequestState.GENERATION_COMPLETE
@@ -564,7 +594,8 @@ def executor_request_to_llm_request(
         cache_salt_id=executor_request.cache_salt_id,
         arrival_time=getattr(executor_request, "py_arrival_time", None),
         py_multimodal_data=getattr(executor_request, "py_multimodal_data",
-                                   None))
+                                   None),
+        top_logprobs=executor_request.output_config.top_logprobs)
     if child_req_ids:
         for child_id in child_req_ids:
             llm_request.create_child_request(child_id)