From 3bb168e2082fcc7e20ca4f9eead983c816b1a9fa Mon Sep 17 00:00:00 2001 From: Stefan Niebler <82932102+stnie@users.noreply.github.com> Date: Wed, 27 Aug 2025 10:08:17 +0000 Subject: [PATCH 01/10] [TRTLLM-5670][feat] Enhance OutputConfig to support top log probabilities - Added a new parameter `topLogProbs` to the `OutputConfig` class, allowing users to specify the number of top log probabilities to return. - Updated related methods and bindings in the executor and API layers to accommodate this feature. - updated code in Torch sampler to provide the functionality - added a new Dict to LogProbStorage containing the top logprobs - included validation checks for top log probabilities in the sampling parameters - added tests to ensure correct functionality. Signed-off-by: Stefan Niebler <82932102+stnie@users.noreply.github.com> --- cpp/include/tensorrt_llm/executor/executor.h | 5 +- cpp/tensorrt_llm/executor/outputConfig.cpp | 3 +- .../nanobind/executor/request.cpp | 8 +- cpp/tensorrt_llm/pybind/executor/request.cpp | 6 +- examples/llm-api/quickstart_advanced.py | 18 +++++ tensorrt_llm/_torch/pyexecutor/_util.py | 6 +- tensorrt_llm/_torch/pyexecutor/config.py | 3 + tensorrt_llm/_torch/pyexecutor/llm_request.py | 57 ++++++++++---- tensorrt_llm/_torch/pyexecutor/sampler.py | 74 +++++++++++++------ tensorrt_llm/executor/result.py | 7 ++ tensorrt_llm/llmapi/llm_args.py | 7 +- tensorrt_llm/sampling_params.py | 16 ++++ .../integration/test_lists/test-db/l0_a30.yml | 1 + .../sampler/test_return_top_k_logprobs.py | 72 ++++++++++++++++++ 14 files changed, 238 insertions(+), 45 deletions(-) create mode 100644 tests/unittest/_torch/sampler/test_return_top_k_logprobs.py diff --git a/cpp/include/tensorrt_llm/executor/executor.h b/cpp/include/tensorrt_llm/executor/executor.h index 9dda07d19c6a..dfaa288d00e9 100644 --- a/cpp/include/tensorrt_llm/executor/executor.h +++ b/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> additionalModelOutputs = std::nullopt); + std::optional> additionalModelOutputs = std::nullopt, + SizeType32 topLogProbs = 0); /// @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> additionalModelOutputs; + /// @brief The number of logprobs to return. Default is 0. + std::optional topLogProbs; }; /// @brief Configuration for speculative decoding with external draft tokens. diff --git a/cpp/tensorrt_llm/executor/outputConfig.cpp b/cpp/tensorrt_llm/executor/outputConfig.cpp index 43fad705cab9..eb7467e85440 100644 --- a/cpp/tensorrt_llm/executor/outputConfig.cpp +++ b/cpp/tensorrt_llm/executor/outputConfig.cpp @@ -22,7 +22,7 @@ namespace tensorrt_llm::executor OutputConfig::OutputConfig(bool inReturnLogProbs, bool inReturnContextLogits, bool inReturnGenerationLogits, bool inExcludeInputFromOutput, bool inReturnEncoderOutput, bool inReturnPerfMetrics, - std::optional> additionalModelOutputs) + std::optional> additionalModelOutputs, SizeType32 topLogProbs) : returnLogProbs(inReturnLogProbs) , returnContextLogits(inReturnContextLogits) , returnGenerationLogits(inReturnGenerationLogits) @@ -30,6 +30,7 @@ OutputConfig::OutputConfig(bool inReturnLogProbs, bool inReturnContextLogits, bo , returnEncoderOutput(inReturnEncoderOutput) , returnPerfMetrics(inReturnPerfMetrics) , additionalModelOutputs(std::move(additionalModelOutputs)) + , topLogProbs(topLogProbs) { } diff --git a/cpp/tensorrt_llm/nanobind/executor/request.cpp b/cpp/tensorrt_llm/nanobind/executor/request.cpp index de9aa8a8c07a..7c5170ab22af 100644 --- a/cpp/tensorrt_llm/nanobind/executor/request.cpp +++ b/cpp/tensorrt_llm/nanobind/executor/request.cpp @@ -216,17 +216,18 @@ void initRequestBindings(nb::module_& m) [](tle::OutputConfig& self, std::optional return_log_probs, std::optional return_context_logits, std::optional return_generation_logits, std::optional exclude_input_from_output, std::optional return_encoder_output, std::optional return_perf_metrics, - std::optional> additional_model_outputs) + std::optional> additional_model_outputs, + std::optional 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.value_or(0)); }, 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) @@ -234,6 +235,7 @@ void initRequestBindings(nb::module_& m) .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); diff --git a/cpp/tensorrt_llm/pybind/executor/request.cpp b/cpp/tensorrt_llm/pybind/executor/request.cpp index 097f598557be..0615b04e0798 100644 --- a/cpp/tensorrt_llm/pybind/executor/request.cpp +++ b/cpp/tensorrt_llm/pybind/executor/request.cpp @@ -204,11 +204,12 @@ void initRequestBindings(pybind11::module_& m) state[6].cast>>()); }; py::class_(m, "OutputConfig") - .def(py::init>>(), + .def(py::init>, + 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") = 0) .def_readwrite("return_log_probs", &tle::OutputConfig::returnLogProbs) .def_readwrite("return_context_logits", &tle::OutputConfig::returnContextLogits) .def_readwrite("return_generation_logits", &tle::OutputConfig::returnGenerationLogits) @@ -216,6 +217,7 @@ void initRequestBindings(pybind11::module_& m) .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) diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py index 61240b496de7..2a7c66fdc033 100644 --- a/examples/llm-api/quickstart_advanced.py +++ b/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=0) + parser.add_argument('--max_top_logprobs', type=int, default=0) return parser @@ -213,6 +215,15 @@ 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.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 == 0) and (args.top_p is None + or args.top_p == 0.0) + + llm = LLM( model=args.model_dir, backend='pytorch', @@ -246,6 +257,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 @@ -265,6 +278,7 @@ def setup_llm(args, **kwargs): return_context_logits=args.return_context_logits, return_generation_logits=args.return_generation_logits, logprobs=args.logprobs, + top_logprobs=args.top_logprobs, n=args.n, best_of=best_of, use_beam_search=use_beam_search) @@ -306,6 +320,10 @@ def main(): ) if args.logprobs: print(f"[{i}]{sequence_id_text} Logprobs: {sequence.logprobs}") + if args.top_logprobs: + print( + f"[{i}]{sequence_id_text} Top logprobs: {sequence.top_logprobs}" + ) if __name__ == '__main__': diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 684158a80641..ea558f36b614 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/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,8 @@ 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) diff --git a/tensorrt_llm/_torch/pyexecutor/config.py b/tensorrt_llm/_torch/pyexecutor/config.py index 7f46c521b6fe..6301a8a15cee 100644 --- a/tensorrt_llm/_torch/pyexecutor/config.py +++ b/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', diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 3d21238ee86f..dd26bafbcf06 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/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], @@ -291,6 +315,7 @@ def __init__( # Detour handling of some parameters client_id: int = None, return_log_probs: bool = False, + top_logprobs: int = 0, return_context_logits: bool = False, return_generation_logits: bool = False, return_logits_device_memory: bool = True, @@ -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 @@ -564,7 +590,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) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py index 1e2897f9adc9..cd5f53e534ad 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler.py @@ -34,6 +34,8 @@ class SampleStateTensors: new_tokens: torch.Tensor log_probs: torch.Tensor | None = None + top_tokens: torch.Tensor | None = None + top_log_probs: torch.Tensor | None = None def values(self): return vars(self).values() @@ -177,7 +179,7 @@ def top_k_sampling_batch(logits, # sample from the distribution and generate result of [batch_size, 1] next_tokens = torch.multinomial(softmax, num_samples=1, generator=generator).squeeze(-1) - return next_tokens, softmax + return next_tokens, softmax, indices def top_p_sampling_batch(logits: torch.Tensor, @@ -214,7 +216,7 @@ def top_p_sampling_batch(logits: torch.Tensor, # sample from the distribution and generate result of [batch_size, 1] next_tokens = torch.multinomial(softmax, num_samples=1, generator=generator).squeeze(-1) - return next_tokens, softmax + return next_tokens, softmax, sorted_indices def top_k_top_p_sampling_batch(logits: torch.Tensor, @@ -260,13 +262,13 @@ def top_k_top_p_sampling_batch(logits: torch.Tensor, # sample from the distribution and generate result of [batch_size, 1] next_tokens = torch.multinomial(softmax, num_samples=1, generator=generator).squeeze(-1) - return next_tokens, softmax + return next_tokens, softmax, indices def greedy_search_sampling_batch(logits): next_tokens = torch.argmax(logits, dim=-1) softmax = torch.softmax(logits, dim=-1) - return next_tokens, softmax + return next_tokens, softmax, None def get_rejected_indices(draft_probs: torch.Tensor, target_probs: torch.Tensor, @@ -414,6 +416,7 @@ class Args: max_draft_len: int max_num_sequences: int max_beam_width: int + max_top_logprobs: int enable_mixed_sampler: bool def __init__(self, args: Args): @@ -454,14 +457,23 @@ def handle_logprobs(self, request: LlmRequest, state: SampleStateTorch, *, current_slice = slice(0, count), request.py_seq_slot, beam if request.py_return_log_probs: assert state.host.log_probs is not None - log_probs = state.host.log_probs[request.py_seq_slot][beam][:count] + log_probs = state.host.log_probs[request.py_seq_slot, beam, :count] current_tokens = state.host.new_tokens[current_slice] - + top_log_probs = state.host.top_log_probs[request.py_seq_slot, + beam, :count] + top_current_tokens = state.host.top_tokens[current_slice] token_log_probs = [{ int(token): Logprob(logprob=logprob, rank=1) } for token, logprob in zip(current_tokens, log_probs.tolist())] + token_top_log_probs = [{ + int(token[k]): + Logprob(logprob=logprob[k], rank=k) + for k in range(request.py_top_logprobs) + } for token, logprob in zip(top_current_tokens, + top_log_probs.tolist())] assert beam == 0, "The following call relies on beam_width to be 1 - hence the list with a single element" - request.py_result.append_log_probs([token_log_probs]) + request.py_result.append_log_probs( + [token_log_probs], top_log_probs=[token_top_log_probs]) FinishReasons: TypeAlias = list[list[int]] """`(num_seq_slots, num_steps)`""" @@ -583,13 +595,13 @@ def update_requests(self, state: SampleStateTorch) -> None: req.py_decoding_iter += 1 def log_probs_host(self, scheduled_requests: ScheduledRequests): - """Shape: In lockstep with TRTLLMSampler: https://github.com/NVIDIA/TensorRT-LLM/blob/cea5dd1e3883b18bf50901a7f196f50a9544c28c/cpp/include/tensorrt_llm/runtime/decoderState.h#L103""" + """Shape: max_num_sequences, max_beam_width == 1, max_tokens, 1 + max_top_logprobs)""" if any(req.py_return_log_probs for req in scheduled_requests.all_requests()): - return torch.empty( - (self.max_num_sequences, SINGLE_BEAM_WIDTH, self.max_tokens), - device="cpu", - pin_memory=True) + return torch.empty((self.max_num_sequences, self.MAX_BEAM_WIDTH, + self.max_tokens, 1 + self.max_top_logprobs), + device="cpu", + pin_memory=True) return None def sample_async( @@ -916,8 +928,8 @@ def _process_requests(self, next_tokens = torch.argmax(logits, dim=-1) self.append_eagle3(next_tokens, model_outputs) int_next_tokens = next_tokens.to(torch.int, non_blocking=True) - next_tokens = int_next_tokens.view(1, -1, SINGLE_BEAM_WIDTH) - new_tokens[:1].index_copy_(1, seq_slots, next_tokens) + next_tokens = int_next_tokens.view(1, -1, SINGLE_BEAM_WIDTH, 1) + new_tokens[:1, ..., :1].index_copy_(1, seq_slots, next_tokens) return strategies = sampling_strategies(requests) @@ -938,13 +950,17 @@ def _process_requests(self, ] logits = self._apply_embedding_bias(logits, requests, steps_per_request) - batched_next_tokens, batched_softmax = sample( + batched_next_tokens, batched_softmax, batched_indices = sample( batched_strategy, logits, generator) self.append_eagle3(batched_next_tokens, model_outputs) offset = 0 for i, (strategy, slot, steps, request) in enumerate( zip(strategies, seq_slots_host, num_steps, requests)): + if request.py_top_logprobs > self.max_top_logprobs: + raise ValueError( + f"top_logprobs {request.py_top_logprobs} cannot exceed max_top_logprobs {self.max_top_logprobs}" + ) input_slice = slice(offset, offset + steps) logits = raw_logits[input_slice] @@ -952,22 +968,38 @@ def _process_requests(self, if batched_next_tokens is None: logits = self._apply_embedding_bias(logits, [req]) - next_tokens, softmax = sample(strategy, logits, generator) + next_tokens, softmax, indices = sample(strategy, logits, + generator) else: # Batched processing already applied bias, just use the results next_tokens = batched_next_tokens[input_slice] softmax = batched_softmax[input_slice] - current_slice = slice(0, steps), slot, BEAM_0 + indices = batched_indices[ + input_slice] if strategy != GREEDY else None + + current_slice = slice(0, steps), slot, BEAM_0, slice( + 0, 1 + request.py_top_logprobs) + if request.py_top_logprobs > 0: + next_tokens = torch.cat([ + next_tokens.unsqueeze(1), + indices[:, :request.py_top_logprobs] + ], + dim=1) + else: + next_tokens = next_tokens.unsqueeze(1) new_tokens[current_slice] = next_tokens if request.py_draft_logits is not None: request.py_target_probs = softmax.clone() if log_probs_host is not None: assert BEAM_0 == 0, "The following call relies on beam_width to be 1 - hence the unsqueeze" - token_probs = torch.gather( - softmax, dim=1, index=next_tokens.unsqueeze(1)).squeeze(-1) + token_probs = torch.gather(softmax, dim=1, + index=next_tokens).reshape( + steps, -1) + log_probs = torch.log(token_probs) - log_probs_host[slot, BEAM_0, :steps].copy_(log_probs, - non_blocking=True) + log_probs_host[slot, BEAM_0, :steps, :1 + + request.py_top_logprobs].copy_(log_probs, + non_blocking=True) offset += steps diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py index ddbe2f636aaf..b31808f4f7c0 100644 --- a/tensorrt_llm/executor/result.py +++ b/tensorrt_llm/executor/result.py @@ -112,6 +112,7 @@ class CompletionOutput: logprobs: Optional[TokenLogprobs | List[float]] = field(default_factory=list) prompt_logprobs: Optional[TokenLogprobs] = field(default_factory=list) + top_logprobs: Optional[TokenLogprobs] = field(default_factory=list) finish_reason: Optional[Literal['stop', 'length', 'timeout', 'cancelled']] = None stop_reason: Optional[Union[int, str]] = None @@ -260,6 +261,12 @@ def _handle_sequence(self, # Therefore, we treat extra logprobs/logits as expected and only consume what's needed. output.logprobs = output.logprobs[:output.length] assert len(output.logprobs) == output.length + if response_tensors.top_log_probs is not None: + output.top_logprobs = response_tensors.top_log_probs[src_idx] + if finish_reasons[src_idx] != tllm.FinishReason.CANCELLED: + if len(output.top_logprobs) > output.length: + output.top_logprobs = output.top_logprobs[:output. + length] if response_tensors.generation_logits is not None: output.generation_logits = response_tensors.generation_logits[ src_idx, :output.length] diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 45929ff60d75..84e2cfae2f42 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -1352,6 +1352,9 @@ class BaseLlmArgs(StrictBaseModel): max_beam_width: Optional[int] = Field(default=None, description="The maximum beam width.") + max_top_logprobs: int = Field( + default=0, description="The maximum number of top logprobs.") + max_num_tokens: Optional[int] = Field( default=None, description="The maximum number of tokens.") @@ -2617,7 +2620,9 @@ def get_pytorch_backend_config(self) -> "PyTorchConfig": attention_dp_batching_wait_iters=self.attention_dp_config. batching_wait_iters if self.attention_dp_config is not None else AttentionDpConfig.model_fields['batching_wait_iters'].default, - batch_wait_timeout_ms=self.batch_wait_timeout_ms) + batch_wait_timeout_ms=self.batch_wait_timeout_ms, + max_top_logprobs=self.max_top_logprobs, + ) def update_llm_args_with_extra_dict( diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index 361c0fc0c0f4..263ed20d604d 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -173,6 +173,7 @@ class SamplingParams: logprobs (int, optional): Number of log probabilities to return per output token. Defaults to None. prompt_logprobs (int, optional): Number of log probabilities to return per prompt token. Defaults to None. + top_logprobs (int, optional): Number of top log probabilities to return per output token. Defaults to 0. return_context_logits (bool): Controls if Result should contain the context logits. Defaults to False. return_generation_logits (bool): Controls if Result should contain the generation logits. Defaults to False. exclude_input_from_output (bool): Controls if output tokens in Result should include the input tokens. Defaults to True. @@ -241,6 +242,7 @@ class SamplingParams: # Keep the below fields in sync with tllme.OutputConfig logprobs: Optional[int] = None prompt_logprobs: Optional[int] = None + top_logprobs: int = 0 return_context_logits: bool = False return_generation_logits: bool = False exclude_input_from_output: bool = True @@ -323,6 +325,20 @@ def _validate(self): self.logprobs = self.logprobs and int(self.logprobs) self.prompt_logprobs = self.prompt_logprobs and int(self.prompt_logprobs) + if self.top_logprobs < 0: + raise ValueError(f"top_logprobs must be >= 0, got {self.top_logprobs}") + + if self._greedy_decoding and self.top_logprobs > 0: + # check for greedy sampling + raise ValueError( + f"top_logprobs {self.top_logprobs} cannot exceed 0 for greedy sampling" + ) + elif self.top_k is not None and self.top_k < self.top_logprobs: + # check for top-k sampling + raise ValueError(f"top_logprobs {self.top_logprobs} cannot exceed top_k {self.top_k}") + if self.top_logprobs > 0 and self.use_beam_search: + raise ValueError("top_logprobs + beam search is not supported") + @property def _greedy_decoding(self) -> bool: return ( diff --git a/tests/integration/test_lists/test-db/l0_a30.yml b/tests/integration/test_lists/test-db/l0_a30.yml index 64cbb936b629..ee4bd4ee2cca 100644 --- a/tests/integration/test_lists/test-db/l0_a30.yml +++ b/tests/integration/test_lists/test-db/l0_a30.yml @@ -21,6 +21,7 @@ l0_a30: - unittest/_torch/modeling -k "modeling_out_of_tree" - unittest/_torch/auto_deploy/unit/singlegpu -k "not test_trtllm_bench_backend_comparison" - unittest/_torch/sampler/test_beam_search.py + - unittest/_torch/sampler/test_return_top_k_logprobs.py - test_e2e.py::test_openai_completions_with_logit_bias[torch_sampler] - test_e2e.py::test_openai_chat_with_logit_bias[torch_sampler] - test_e2e.py::test_openai_completions_with_logit_bias[trtllm_sampler] diff --git a/tests/unittest/_torch/sampler/test_return_top_k_logprobs.py b/tests/unittest/_torch/sampler/test_return_top_k_logprobs.py new file mode 100644 index 000000000000..62e1190018a4 --- /dev/null +++ b/tests/unittest/_torch/sampler/test_return_top_k_logprobs.py @@ -0,0 +1,72 @@ +import os + +import pytest +from utils.llm_data import llm_models_root +from utils.util import force_ampere + +from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm.llmapi import CudaGraphConfig, KvCacheConfig + + +@pytest.fixture(scope="module") +def input_prompts(): + return [ + "Born in north-east France, Soyer trained as a", + "The future of AI is", + ] + + +@pytest.fixture(scope="module") +def llm_torch(input_prompts): + return LLM( + model=os.path.join(llm_models_root(), "llama-models-v2", + "TinyLlama-1.1B-Chat-v1.0"), + kv_cache_config=KvCacheConfig(max_tokens=10000), + max_batch_size=len( + input_prompts + ), # use small batch size to prevent large buffers from possibly hiding wrong data accesses. + max_seq_len=32, + disable_overlap_scheduler=False, + sampler_type="TorchSampler", + cuda_graph_config=CudaGraphConfig(batch_sizes=[1, 2, 4, 8], + enable_padding=True), + enable_mixed_sampler=True, + max_top_logprobs=4, + ) + + +@force_ampere # Save H100 resource +@pytest.mark.parametrize("top_logprobs", [0, 1, 2, 4]) +@pytest.mark.parametrize("top_k", [None, 2]) +@pytest.mark.parametrize("top_p", [None, 0.5]) +@pytest.mark.threadleak(enabled=False) +def test_generate_with_top_logprobs(top_logprobs: int, top_k: int, top_p: float, + llm_torch, input_prompts): + max_tokens = 8 + is_topk = top_k is not None and top_k > 1 + is_topp = top_p is not None and top_p > 0.0 and not is_topk + is_valid_setup = is_topk and top_logprobs <= top_k or is_topp or top_logprobs == 0 + if is_valid_setup: + # passing testcases + sampling_params = SamplingParams(max_tokens=max_tokens, + logprobs=True, + top_k=top_k, + top_p=top_p, + temperature=1.0, + top_logprobs=top_logprobs) + for output in llm_torch.generate(input_prompts, + sampling_params=sampling_params): + if top_logprobs > 0: + assert len(output.outputs[0].logprobs) == max_tokens + assert len(output.outputs[0].top_logprobs) == max_tokens + assert len(output.outputs[0].logprobs[0]) == 1 + assert len(output.outputs[0].top_logprobs[0]) == top_logprobs + else: + # expected to fail testcases + with pytest.raises(ValueError, match="top_logprobs.*cannot exceed.*"): + sampling_params = SamplingParams(max_tokens=max_tokens, + logprobs=True, + top_k=top_k, + top_p=top_p, + temperature=1.0, + top_logprobs=top_logprobs) From e510ca05b38f2e012e7b55d0fd3bcee275b61e34 Mon Sep 17 00:00:00 2001 From: Stefan Niebler <82932102+stnie@users.noreply.github.com> Date: Wed, 27 Aug 2025 16:16:15 +0000 Subject: [PATCH 02/10] [TRTLLM-5670][chore] Update OutputConfig and related components for optional top log probabilities - Changed `topLogProbs` parameter in `OutputConfig` to be optional, allowing for more flexible configurations. - Updated related methods and bindings in the executor and API layers to handle the new optional type. - Enhanced validation checks in sampling parameters to ensure correct usage of `top_logprobs`. - Adjusted tests to accommodate the changes in parameter handling and validation. Signed-off-by: Stefan Niebler <82932102+stnie@users.noreply.github.com> --- cpp/include/tensorrt_llm/executor/executor.h | 2 +- cpp/tensorrt_llm/executor/outputConfig.cpp | 2 +- .../nanobind/executor/request.cpp | 10 ++- cpp/tensorrt_llm/pybind/executor/request.cpp | 10 ++- examples/llm-api/quickstart_advanced.py | 4 +- tensorrt_llm/_torch/pyexecutor/llm_request.py | 6 +- tensorrt_llm/_torch/pyexecutor/sampler.py | 84 +++++++++++++++---- tensorrt_llm/sampling_params.py | 34 ++++---- .../sampler/test_return_top_k_logprobs.py | 32 ++++++- .../bindings/test_executor_bindings.py | 7 +- 10 files changed, 143 insertions(+), 48 deletions(-) diff --git a/cpp/include/tensorrt_llm/executor/executor.h b/cpp/include/tensorrt_llm/executor/executor.h index dfaa288d00e9..6bc0b5f247d5 100644 --- a/cpp/include/tensorrt_llm/executor/executor.h +++ b/cpp/include/tensorrt_llm/executor/executor.h @@ -216,7 +216,7 @@ class OutputConfig bool returnGenerationLogits = false, bool excludeInputFromOutput = false, bool returnEncoderOutput = false, bool returnPerfMetrics = false, std::optional> additionalModelOutputs = std::nullopt, - SizeType32 topLogProbs = 0); + std::optional topLogProbs = std::nullopt); /// @brief Controls if Result should contain log probabilities. Default is false. bool returnLogProbs; diff --git a/cpp/tensorrt_llm/executor/outputConfig.cpp b/cpp/tensorrt_llm/executor/outputConfig.cpp index eb7467e85440..8526a3878218 100644 --- a/cpp/tensorrt_llm/executor/outputConfig.cpp +++ b/cpp/tensorrt_llm/executor/outputConfig.cpp @@ -22,7 +22,7 @@ namespace tensorrt_llm::executor OutputConfig::OutputConfig(bool inReturnLogProbs, bool inReturnContextLogits, bool inReturnGenerationLogits, bool inExcludeInputFromOutput, bool inReturnEncoderOutput, bool inReturnPerfMetrics, - std::optional> additionalModelOutputs, SizeType32 topLogProbs) + std::optional> additionalModelOutputs, std::optional topLogProbs) : returnLogProbs(inReturnLogProbs) , returnContextLogits(inReturnContextLogits) , returnGenerationLogits(inReturnGenerationLogits) diff --git a/cpp/tensorrt_llm/nanobind/executor/request.cpp b/cpp/tensorrt_llm/nanobind/executor/request.cpp index 7c5170ab22af..d6fe21c85026 100644 --- a/cpp/tensorrt_llm/nanobind/executor/request.cpp +++ b/cpp/tensorrt_llm/nanobind/executor/request.cpp @@ -198,17 +198,19 @@ 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(state[0]), nb::cast(state[1]), nb::cast(state[2]), nb::cast(state[3]), nb::cast(state[4]), nb::cast(state[5]), - nb::cast>>(state[6])); + nb::cast>>(state[6]), + nb::cast>(state[7])); }; nb::class_(m, "OutputConfig") .def( @@ -222,7 +224,7 @@ void initRequestBindings(nb::module_& m) 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, top_logprobs.value_or(0)); + 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(), diff --git a/cpp/tensorrt_llm/pybind/executor/request.cpp b/cpp/tensorrt_llm/pybind/executor/request.cpp index 0615b04e0798..8fcc46c6e8cc 100644 --- a/cpp/tensorrt_llm/pybind/executor/request.cpp +++ b/cpp/tensorrt_llm/pybind/executor/request.cpp @@ -191,17 +191,19 @@ 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(), state[1].cast(), state[2].cast(), state[3].cast(), state[4].cast(), state[5].cast(), - state[6].cast>>()); + state[6].cast>>(), + state[7].cast>()); }; py::class_(m, "OutputConfig") .def(py::init>, @@ -209,7 +211,7 @@ void initRequestBindings(pybind11::module_& m) 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("top_logprobs") = 0) + 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) diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py index 2a7c66fdc033..b0212a048a37 100644 --- a/examples/llm-api/quickstart_advanced.py +++ b/examples/llm-api/quickstart_advanced.py @@ -147,7 +147,7 @@ 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=0) + parser.add_argument('--top_logprobs', type=int, default=None) parser.add_argument('--max_top_logprobs', type=int, default=0) return parser @@ -216,7 +216,7 @@ def setup_llm(args, **kwargs): ) # if only top_logprobs is set, set max_top_logprobs to top_logprobs - if args.max_top_logprobs < args.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 ( diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index dd26bafbcf06..4ce64c2ea1c3 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -315,7 +315,6 @@ def __init__( # Detour handling of some parameters client_id: int = None, return_log_probs: bool = False, - top_logprobs: int = 0, return_context_logits: bool = False, return_generation_logits: bool = False, return_logits_device_memory: bool = True, @@ -327,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", @@ -420,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 + def finish_by(self, reason: FinishReason, beam: int) -> None: """CPP finish by reason does not support beam_width > 1""" self.state = LlmRequestState.GENERATION_COMPLETE diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py index cd5f53e534ad..08c8bb1d4e29 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler.py @@ -158,6 +158,12 @@ def is_generation_model(self) -> bool: def top_k_sampling_batch(logits, top_k=50, generator: Optional[torch.Generator] = None): + """Sample tokens using top_k sampling. + Returns: + next_tokens: [batch_size, 1] (sampled tokens) + softmax: [batch_size, vocab_size] (probability distribution) + indices: [batch_size, top_k] (indices of top_k logits) + """ logits_dim = logits.dim() if logits_dim == 1: logits = logits.unsqueeze(0) @@ -186,6 +192,12 @@ def top_p_sampling_batch(logits: torch.Tensor, top_p: float = 0.9, temperature: float = 1.0, generator: Optional[torch.Generator] = None): + """Sample tokens using top_p sampling. + Returns: + next_tokens: [batch_size, 1] (sampled tokens) + softmax: [batch_size, vocab_size] (probability distribution) + indices: [batch_size, top_k] (indices of top_k logits) + """ logits_dim = logits.dim() if logits_dim == 1: logits = logits.unsqueeze(0) @@ -224,6 +236,12 @@ def top_k_top_p_sampling_batch(logits: torch.Tensor, top_p: float, temperature: float = 1.0, generator: Optional[torch.Generator] = None): + """Sample tokens using top_k_top_p sampling. + Returns: + next_tokens: [batch_size, 1] (sampled tokens) + softmax: [batch_size, vocab_size] (probability distribution) + indices: [batch_size, top_k] (indices of top_k logits) + """ logits_dim = logits.dim() if logits_dim == 1: logits = logits.unsqueeze(0) @@ -266,6 +284,12 @@ def top_k_top_p_sampling_batch(logits: torch.Tensor, def greedy_search_sampling_batch(logits): + """Sample tokens using greedy search. + Returns: + next_tokens: [batch_size, 1] (sampled tokens) + softmax: [batch_size, vocab_size] (probability distribution) + indices: None (greedy search does not return indices) + """ next_tokens = torch.argmax(logits, dim=-1) softmax = torch.softmax(logits, dim=-1) return next_tokens, softmax, None @@ -437,6 +461,10 @@ def __init__(self, args: Args): self._global_seed = 42 self._generator = None + @property + def calculate_top_logprobs(self): + return self.max_top_logprobs > 0 + def get_generator(self, device: torch.device) -> torch.Generator: """Get a deterministic generator for the specified device. @@ -462,15 +490,30 @@ def handle_logprobs(self, request: LlmRequest, state: SampleStateTorch, *, top_log_probs = state.host.top_log_probs[request.py_seq_slot, beam, :count] top_current_tokens = state.host.top_tokens[current_slice] - token_log_probs = [{ - int(token): Logprob(logprob=logprob, rank=1) - } for token, logprob in zip(current_tokens, log_probs.tolist())] token_top_log_probs = [{ int(token[k]): Logprob(logprob=logprob[k], rank=k) for k in range(request.py_top_logprobs) } for token, logprob in zip(top_current_tokens, top_log_probs.tolist())] + return [token_top_log_probs] + else: + return None + + def handle_logprobs(self, request: LlmRequest, state: SampleState, *, + beam: int, count: int): + current_slice = slice(0, count), request.py_seq_slot, beam + if request.py_return_log_probs: + assert state.host.log_probs is not None + log_probs = state.host.log_probs[request.py_seq_slot, beam, :count] + current_tokens = state.host.new_tokens[current_slice] + token_log_probs = [{ + int(token): Logprob(logprob=logprob, rank=1) + } for token, logprob in zip(current_tokens, log_probs.tolist())] + token_top_log_probs = self._maybe_get_top_logprobs(request, + state, + beam=beam, + count=count) assert beam == 0, "The following call relies on beam_width to be 1 - hence the list with a single element" request.py_result.append_log_probs( [token_log_probs], top_log_probs=[token_top_log_probs]) @@ -595,7 +638,8 @@ def update_requests(self, state: SampleStateTorch) -> None: req.py_decoding_iter += 1 def log_probs_host(self, scheduled_requests: ScheduledRequests): - """Shape: max_num_sequences, max_beam_width == 1, max_tokens, 1 + max_top_logprobs)""" + """Shape: max_num_sequences, max_beam_width == 1, max_tokens, 1 + max_top_logprobs) + The last dimension contains log_probs for the sampled tokens and additionally log_probs for the top tokens if top_logprobs is specified""" if any(req.py_return_log_probs for req in scheduled_requests.all_requests()): return torch.empty((self.max_num_sequences, self.MAX_BEAM_WIDTH, @@ -631,6 +675,15 @@ def sample_async( new_tokens_host = new_tokens.to(device="cpu", non_blocking=True) finish_reasons_host = finish_reasons.to(device="cpu", non_blocking=True) + # Setup these optional values if possible + log_probs = None + top_log_probs = None + top_tokens = None + if log_probs_host is not None: + log_probs = log_probs_host[..., 0] + if self.calculate_top_logprobs: + top_log_probs = log_probs_host[..., 1:] + top_tokens = new_tokens[..., 1:] sampler_event = torch.cuda.Event() sampler_event.record() return SampleStateTorch( @@ -957,7 +1010,10 @@ def _process_requests(self, offset = 0 for i, (strategy, slot, steps, request) in enumerate( zip(strategies, seq_slots_host, num_steps, requests)): - if request.py_top_logprobs > self.max_top_logprobs: + # if request.calculate_top_logprobs, then request.py_top_logprobs is the number of top tokens to sample + # otherwise, it is None and should be defaulted to0 + num_top_tokens = request.py_top_logprobs if request.calculate_top_logprobs else 0 + if num_top_tokens > self.max_top_logprobs: raise ValueError( f"top_logprobs {request.py_top_logprobs} cannot exceed max_top_logprobs {self.max_top_logprobs}" ) @@ -978,13 +1034,12 @@ def _process_requests(self, input_slice] if strategy != GREEDY else None current_slice = slice(0, steps), slot, BEAM_0, slice( - 0, 1 + request.py_top_logprobs) - if request.py_top_logprobs > 0: - next_tokens = torch.cat([ - next_tokens.unsqueeze(1), - indices[:, :request.py_top_logprobs] - ], - dim=1) + 0, 1 + num_top_tokens) + if request.calculate_top_logprobs and strategy != GREEDY: + # concatenate the sampled tokens with the top tokens + next_tokens = torch.cat( + [next_tokens.unsqueeze(1), indices[:, :num_top_tokens]], + dim=1) else: next_tokens = next_tokens.unsqueeze(1) new_tokens[current_slice] = next_tokens @@ -997,9 +1052,8 @@ def _process_requests(self, steps, -1) log_probs = torch.log(token_probs) - log_probs_host[slot, BEAM_0, :steps, :1 + - request.py_top_logprobs].copy_(log_probs, - non_blocking=True) + log_probs_host[slot, BEAM_0, :steps, :1 + num_top_tokens].copy_( + log_probs, non_blocking=True) offset += steps diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index 263ed20d604d..52413db7e059 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -242,7 +242,7 @@ class SamplingParams: # Keep the below fields in sync with tllme.OutputConfig logprobs: Optional[int] = None prompt_logprobs: Optional[int] = None - top_logprobs: int = 0 + top_logprobs: Optional[int] = None return_context_logits: bool = False return_generation_logits: bool = False exclude_input_from_output: bool = True @@ -325,19 +325,25 @@ def _validate(self): self.logprobs = self.logprobs and int(self.logprobs) self.prompt_logprobs = self.prompt_logprobs and int(self.prompt_logprobs) - if self.top_logprobs < 0: - raise ValueError(f"top_logprobs must be >= 0, got {self.top_logprobs}") - - if self._greedy_decoding and self.top_logprobs > 0: - # check for greedy sampling - raise ValueError( - f"top_logprobs {self.top_logprobs} cannot exceed 0 for greedy sampling" - ) - elif self.top_k is not None and self.top_k < self.top_logprobs: - # check for top-k sampling - raise ValueError(f"top_logprobs {self.top_logprobs} cannot exceed top_k {self.top_k}") - if self.top_logprobs > 0 and self.use_beam_search: - raise ValueError("top_logprobs + beam search is not supported") + if self.top_logprobs is not None: + if not self.logprobs: + raise ValueError("You need to set logprobs to True to use top_logprobs.") + if self.top_logprobs < 1: + raise ValueError(f"top_logprobs must be >= 1, got {self.top_logprobs}") + + if self._greedy_decoding: + # check for greedy sampling + raise ValueError( + "top_logprobs must not be specified for greedy sampling, " + "as it provides the same output as logprobs in this case." + ) + elif self.top_k is not None and self.top_k < self.top_logprobs: + # check for top-k sampling + raise ValueError( + f"top_logprobs {self.top_logprobs} must not exceed top_k {self.top_k}" + ) + if self.use_beam_search: + raise ValueError("top_logprobs + beam search is not supported") @property def _greedy_decoding(self) -> bool: diff --git a/tests/unittest/_torch/sampler/test_return_top_k_logprobs.py b/tests/unittest/_torch/sampler/test_return_top_k_logprobs.py index 62e1190018a4..31a4ba9e9a0f 100644 --- a/tests/unittest/_torch/sampler/test_return_top_k_logprobs.py +++ b/tests/unittest/_torch/sampler/test_return_top_k_logprobs.py @@ -36,7 +36,7 @@ def llm_torch(input_prompts): @force_ampere # Save H100 resource -@pytest.mark.parametrize("top_logprobs", [0, 1, 2, 4]) +@pytest.mark.parametrize("top_logprobs", [None, 0, 2, 4]) @pytest.mark.parametrize("top_k", [None, 2]) @pytest.mark.parametrize("top_p", [None, 0.5]) @pytest.mark.threadleak(enabled=False) @@ -45,7 +45,10 @@ def test_generate_with_top_logprobs(top_logprobs: int, top_k: int, top_p: float, max_tokens = 8 is_topk = top_k is not None and top_k > 1 is_topp = top_p is not None and top_p > 0.0 and not is_topk - is_valid_setup = is_topk and top_logprobs <= top_k or is_topp or top_logprobs == 0 + uses_top_logprobs = top_logprobs is not None + is_valid_top_logprobs = top_logprobs > 0 if uses_top_logprobs else True + is_valid_setup = not uses_top_logprobs or (is_valid_top_logprobs and ( + (is_topk and top_logprobs <= top_k) or is_topp)) if is_valid_setup: # passing testcases sampling_params = SamplingParams(max_tokens=max_tokens, @@ -56,17 +59,38 @@ def test_generate_with_top_logprobs(top_logprobs: int, top_k: int, top_p: float, top_logprobs=top_logprobs) for output in llm_torch.generate(input_prompts, sampling_params=sampling_params): - if top_logprobs > 0: + if top_logprobs is not None and top_logprobs > 0: assert len(output.outputs[0].logprobs) == max_tokens assert len(output.outputs[0].top_logprobs) == max_tokens assert len(output.outputs[0].logprobs[0]) == 1 assert len(output.outputs[0].top_logprobs[0]) == top_logprobs else: # expected to fail testcases - with pytest.raises(ValueError, match="top_logprobs.*cannot exceed.*"): + with pytest.raises(ValueError, match="top_logprobs.*"): sampling_params = SamplingParams(max_tokens=max_tokens, logprobs=True, top_k=top_k, top_p=top_p, temperature=1.0, top_logprobs=top_logprobs) + + +@force_ampere # Save H100 resource +@pytest.mark.threadleak(enabled=False) +def test_generate_with_top_logprobs_and_disabler_logprobs( + llm_torch, input_prompts): + max_tokens = 8 + top_logprobs = 2 + top_k = 2 + top_p = None + + # expected to fail testcases + with pytest.raises( + ValueError, + match=".*You need to set logprobs to True to use top_logprobs.*"): + sampling_params = SamplingParams(max_tokens=max_tokens, + logprobs=False, + top_k=top_k, + top_p=top_p, + temperature=1.0, + top_logprobs=top_logprobs) diff --git a/tests/unittest/bindings/test_executor_bindings.py b/tests/unittest/bindings/test_executor_bindings.py index 8556cf54d69b..316e57317d15 100644 --- a/tests/unittest/bindings/test_executor_bindings.py +++ b/tests/unittest/bindings/test_executor_bindings.py @@ -889,10 +889,11 @@ def test_output_config(): assert config.return_encoder_output == False assert config.return_perf_metrics == False assert config.additional_model_outputs is None + assert config.top_logprobs is None config = trtllm.OutputConfig( True, False, True, False, True, False, - list([trtllm.AdditionalModelOutput("topKLogits", True)])) + list([trtllm.AdditionalModelOutput("topKLogits", True)]), 10) assert config.return_log_probs == True assert config.return_context_logits == False assert config.return_generation_logits == True @@ -903,12 +904,13 @@ def test_output_config(): additional_model_output = config.additional_model_outputs[0] assert additional_model_output.name == "topKLogits" assert additional_model_output.gather_context == True + assert config.top_logprobs == 10 def test_output_config_pickle(): config = trtllm.OutputConfig( True, False, True, False, True, False, - list([trtllm.AdditionalModelOutput("topKLogits", True)])) + list([trtllm.AdditionalModelOutput("topKLogits", True)]), 10) config_copy = pickle.loads(pickle.dumps(config)) assert config_copy.return_log_probs == True assert config_copy.return_context_logits == False @@ -920,6 +922,7 @@ def test_output_config_pickle(): additional_model_output = config_copy.additional_model_outputs[0] assert additional_model_output.name == "topKLogits" assert additional_model_output.gather_context == True + assert config_copy.top_logprobs == 10 def test_external_draft_tokens_config(): From 0aeb5d69ac54565eab12cc63712d7338085ed976 Mon Sep 17 00:00:00 2001 From: Stefan Niebler <82932102+stnie@users.noreply.github.com> Date: Thu, 28 Aug 2025 07:24:59 +0000 Subject: [PATCH 03/10] [TRTLLM-5670][chore] Refine several small bugs / descriptions - Updated `OutputConfig` to use `std::optional` for the `top_logprobs` parameter - Adjusted docstring in `sampling_params.py` to correctly state the default of `top_logprobs` as None. - Modified `quickstart_advanced.py` to ensure correct handling of `top_k` conditions. - Clarified documentation in `sampler.py` regarding indices after applying top-k and top-p masks. Signed-off-by: Stefan Niebler <82932102+stnie@users.noreply.github.com> --- cpp/tensorrt_llm/pybind/executor/request.cpp | 2 +- examples/llm-api/quickstart_advanced.py | 2 +- tensorrt_llm/_torch/pyexecutor/sampler.py | 2 +- tensorrt_llm/sampling_params.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/cpp/tensorrt_llm/pybind/executor/request.cpp b/cpp/tensorrt_llm/pybind/executor/request.cpp index 8fcc46c6e8cc..55ae7fb503a5 100644 --- a/cpp/tensorrt_llm/pybind/executor/request.cpp +++ b/cpp/tensorrt_llm/pybind/executor/request.cpp @@ -207,7 +207,7 @@ void initRequestBindings(pybind11::module_& m) }; py::class_(m, "OutputConfig") .def(py::init>, - SizeType32>(), + std::optional>(), 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, diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py index b0212a048a37..40fb065f644d 100644 --- a/examples/llm-api/quickstart_advanced.py +++ b/examples/llm-api/quickstart_advanced.py @@ -220,7 +220,7 @@ def setup_llm(args, **kwargs): 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 == 0) and (args.top_p is None + args.top_k is None or args.top_k == 1) and (args.top_p is None or args.top_p == 0.0) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py index 08c8bb1d4e29..07db4526a31b 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler.py @@ -240,7 +240,7 @@ def top_k_top_p_sampling_batch(logits: torch.Tensor, Returns: next_tokens: [batch_size, 1] (sampled tokens) softmax: [batch_size, vocab_size] (probability distribution) - indices: [batch_size, top_k] (indices of top_k logits) + indices: [batch_size, top_k] (indices after applying top-k then top-p mask) """ logits_dim = logits.dim() if logits_dim == 1: diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index 52413db7e059..ac635bd63c75 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -173,7 +173,7 @@ class SamplingParams: logprobs (int, optional): Number of log probabilities to return per output token. Defaults to None. prompt_logprobs (int, optional): Number of log probabilities to return per prompt token. Defaults to None. - top_logprobs (int, optional): Number of top log probabilities to return per output token. Defaults to 0. + top_logprobs (int, optional): Number of top log probabilities to return per output token. Defaults to None. return_context_logits (bool): Controls if Result should contain the context logits. Defaults to False. return_generation_logits (bool): Controls if Result should contain the generation logits. Defaults to False. exclude_input_from_output (bool): Controls if output tokens in Result should include the input tokens. Defaults to True. From 5022c6bdfa16843eef186302f0785ab2dbaff307 Mon Sep 17 00:00:00 2001 From: Stefan Niebler <82932102+stnie@users.noreply.github.com> Date: Thu, 28 Aug 2025 11:11:26 +0000 Subject: [PATCH 04/10] [TRTLLM-5670][chore] Update autodeploy code to adjust to the top_logprobs changes - Correctly set max_top_logprobs parameter in create_autodeploy_executor. - Updated DemoEngine to correctly unpack top_k and greedy_search Signed-off-by: Stefan Niebler <82932102+stnie@users.noreply.github.com> --- tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py | 1 + tensorrt_llm/_torch/auto_deploy/shim/demollm.py | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py b/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py index d7e82b983843..5e1a86a25365 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py +++ b/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, ) diff --git a/tensorrt_llm/_torch/auto_deploy/shim/demollm.py b/tensorrt_llm/_torch/auto_deploy/shim/demollm.py index 6f93e2476589..c832809e182e 100644 --- a/tensorrt_llm/_torch/auto_deploy/shim/demollm.py +++ b/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 From 9526f1d0e2ed274c281d3f2cec2903adb4e65877 Mon Sep 17 00:00:00 2001 From: Stefan Niebler <82932102+stnie@users.noreply.github.com> Date: Mon, 1 Sep 2025 14:13:36 +0000 Subject: [PATCH 05/10] [TRTLLM-5670][chore] Add top_logprobs parameter to CompletionOutput docstring and updated yaml files for api stability testing - Updated the test cases in `test_llm_api.py`, `completion_output.yaml`, `llm.yaml`, and `sampling_params.yaml` to include the new `top_logprobs` parameter and its associated properties. Signed-off-by: Stefan Niebler <82932102+stnie@users.noreply.github.com> --- tensorrt_llm/executor/result.py | 1 + .../unittest/api_stability/references/completion_output.yaml | 3 +++ tests/unittest/api_stability/references/llm.yaml | 4 ++++ tests/unittest/api_stability/references/sampling_params.yaml | 3 +++ tests/unittest/api_stability/test_llm_api.py | 2 +- 5 files changed, 12 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py index b31808f4f7c0..6f3eadcf614e 100644 --- a/tensorrt_llm/executor/result.py +++ b/tensorrt_llm/executor/result.py @@ -93,6 +93,7 @@ class CompletionOutput: cumulative_logprob (float, optional): The cumulative log probability of the generated output text. Defaults to None. logprobs (TokenLogprobs | List[float], optional): The log probabilities of the top probability words at each position if the logprobs are requested. Defaults to None. prompt_logprobs (TokenLogprobs, optional): The log probabilities per prompt token. Defaults to None. + top_logprobs (TokenLogprobs, optional): The log probabilities of the top probability words at each position if the top_logprobs are requested. Defaults to None. 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. diff --git a/tests/unittest/api_stability/references/completion_output.yaml b/tests/unittest/api_stability/references/completion_output.yaml index cd1f49a89cf4..06159b2446eb 100644 --- a/tests/unittest/api_stability/references/completion_output.yaml +++ b/tests/unittest/api_stability/references/completion_output.yaml @@ -10,6 +10,9 @@ methods: request_perf_metrics: annotation: Optional[tensorrt_llm.bindings.executor.RequestPerfMetrics] default: null + top_logprobs: + annotation: Optional[list[dict[int, tensorrt_llm.executor.result.Logprob]]] + default: null return_annotation: None properties: length: diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 51518bbccdb1..320defa41a24 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -175,6 +175,10 @@ methods: annotation: bool default: False status: prototype + max_top_logprobs: + annotation: int + default: 0 + status: prototype return_annotation: None generate: parameters: diff --git a/tests/unittest/api_stability/references/sampling_params.yaml b/tests/unittest/api_stability/references/sampling_params.yaml index 1ac178ed755a..f460f264a0e9 100644 --- a/tests/unittest/api_stability/references/sampling_params.yaml +++ b/tests/unittest/api_stability/references/sampling_params.yaml @@ -11,5 +11,8 @@ methods: beam_width_array: annotation: Optional[List[int]] default: null + top_logprobs: + annotation: Optional[int] + default: null return_annotation: None properties: {} diff --git a/tests/unittest/api_stability/test_llm_api.py b/tests/unittest/api_stability/test_llm_api.py index 6960f993286c..2b8b5af32971 100644 --- a/tests/unittest/api_stability/test_llm_api.py +++ b/tests/unittest/api_stability/test_llm_api.py @@ -56,7 +56,7 @@ def test_get_output_config(self): "return_log_probs", "return_context_logits", "return_generation_logits", "exclude_input_from_output", "return_encoder_output", "return_perf_metrics", - "additional_model_outputs" + "additional_model_outputs", "top_logprobs" } found_fields = { f From 51175be881dfbd936ea4e53fca8c34d9d2fd1d76 Mon Sep 17 00:00:00 2001 From: Stefan Niebler <82932102+stnie@users.noreply.github.com> Date: Mon, 8 Sep 2025 15:51:10 +0000 Subject: [PATCH 06/10] [TRTLLM-5670][chore] Update top logprobs to reflect changes in sampler - Updated `TorchStore` and `TorchSampler` to include `max_top_logprobs` parameter for managing top log probabilities. - Modified methods to handle top log probabilities during sampling and logging. - Added new Tensors to `TorchStore` containing data for top logprobs - Adjusted tests to validate the correct functionality of top log probabilities in various scenarios. Signed-off-by: Stefan Niebler <82932102+stnie@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/llm_request.py | 2 +- tensorrt_llm/_torch/pyexecutor/sampler.py | 132 ++++++++++-------- tensorrt_llm/_torch/speculative/mtp.py | 8 +- .../sampler/test_return_top_k_logprobs.py | 11 +- 4 files changed, 90 insertions(+), 63 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 4ce64c2ea1c3..6078322229d5 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -422,7 +422,7 @@ def is_dummy(self): @property def calculate_top_logprobs(self): - return self.py_top_logprobs is not None + 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""" diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py index 07db4526a31b..b52e1d24080c 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler.py @@ -388,11 +388,12 @@ def int_tensor(shape: tuple[int, ...], device: str = 'cuda') -> torch.Tensor: class TorchStore: def __init__(self, *, max_draft_len: int, max_num_sequences: int, - max_beam_width: int): + max_beam_width: int, max_top_logprobs: int): self.max_draft_len = max_draft_len self.max_num_sequences = max_num_sequences self.max_beam_width = max_beam_width self.max_tokens = max_draft_len + 1 + self.max_top_logprobs = max_top_logprobs assert max_beam_width == SINGLE_BEAM_WIDTH, "TorchSampler only supports beam_width = 1" self.new_tokens = int_tensor( (self.max_tokens, max_num_sequences, max_beam_width)) @@ -416,6 +417,19 @@ def __init__(self, *, max_draft_len: int, max_num_sequences: int, FinishReason.TIMED_OUT, FinishReason.CANCELLED ] # `in FinishReason` clashes with PyBind11: `TypeError: 'pybind11_type' object is not iterable` } + """Preallocate buffer needed to provide top logprobs""" + if self.max_top_logprobs > 0: + self.top_log_probs_host = torch.empty( + (max_num_sequences, max_beam_width, self.max_tokens, + max_top_logprobs), + device="cpu", + pin_memory=True) + self.top_tokens_host = int_tensor( + (self.max_tokens, max_num_sequences, max_beam_width, + max_top_logprobs)) + else: + self.top_log_probs_host = None + self.top_tokens_host = None @dataclass(kw_only=True) @@ -446,6 +460,7 @@ class Args: def __init__(self, args: Args): self.max_seq_len = args.max_seq_len self.enable_mixed_sampler = args.enable_mixed_sampler + self.max_top_logprobs = args.max_top_logprobs # AutoDeploy build creates the sampler in inference mode, # which would disallow in-place mutating of new_tokens. @@ -453,7 +468,8 @@ def __init__(self, args: Args): with torch.inference_mode(False): self.store = TorchStore(max_draft_len=args.max_draft_len, max_num_sequences=args.max_num_sequences, - max_beam_width=args.max_beam_width) + max_beam_width=args.max_beam_width, + max_top_logprobs=args.max_top_logprobs) self.max_num_sequences = args.max_num_sequences self.max_tokens = self.store.max_tokens @@ -461,10 +477,6 @@ def __init__(self, args: Args): self._global_seed = 42 self._generator = None - @property - def calculate_top_logprobs(self): - return self.max_top_logprobs > 0 - def get_generator(self, device: torch.device) -> torch.Generator: """Get a deterministic generator for the specified device. @@ -480,19 +492,18 @@ def get_generator(self, device: torch.device) -> torch.Generator: self._generator.manual_seed(self._global_seed) return self._generator - def handle_logprobs(self, request: LlmRequest, state: SampleStateTorch, *, - beam: int, count: int): + def _maybe_get_top_logprobs(self, request: LlmRequest, + state: SampleStateTorch, *, beam: int, + count: int): current_slice = slice(0, count), request.py_seq_slot, beam - if request.py_return_log_probs: + if request.calculate_top_logprobs: assert state.host.log_probs is not None - log_probs = state.host.log_probs[request.py_seq_slot, beam, :count] - current_tokens = state.host.new_tokens[current_slice] top_log_probs = state.host.top_log_probs[request.py_seq_slot, beam, :count] top_current_tokens = state.host.top_tokens[current_slice] token_top_log_probs = [{ int(token[k]): - Logprob(logprob=logprob[k], rank=k) + Logprob(logprob=logprob[k], rank=1) for k in range(request.py_top_logprobs) } for token, logprob in zip(top_current_tokens, top_log_probs.tolist())] @@ -516,7 +527,7 @@ def handle_logprobs(self, request: LlmRequest, state: SampleState, *, count=count) assert beam == 0, "The following call relies on beam_width to be 1 - hence the list with a single element" request.py_result.append_log_probs( - [token_log_probs], top_log_probs=[token_top_log_probs]) + [token_log_probs], top_log_probs=token_top_log_probs) FinishReasons: TypeAlias = list[list[int]] """`(num_seq_slots, num_steps)`""" @@ -637,15 +648,15 @@ def update_requests(self, state: SampleStateTorch) -> None: self.handle_logprobs(req, state, beam=BEAM_0, count=processed) req.py_decoding_iter += 1 - def log_probs_host(self, scheduled_requests: ScheduledRequests): - """Shape: max_num_sequences, max_beam_width == 1, max_tokens, 1 + max_top_logprobs) - The last dimension contains log_probs for the sampled tokens and additionally log_probs for the top tokens if top_logprobs is specified""" + def log_probs_host( + self, scheduled_requests: ScheduledRequests) -> torch.Tensor | None: + """Shape: (max_num_sequences, max_beam_width == 1, max_tokens)""" if any(req.py_return_log_probs for req in scheduled_requests.all_requests()): - return torch.empty((self.max_num_sequences, self.MAX_BEAM_WIDTH, - self.max_tokens, 1 + self.max_top_logprobs), - device="cpu", - pin_memory=True) + return torch.empty( + (self.max_num_sequences, SINGLE_BEAM_WIDTH, self.max_tokens), + device="cpu", + pin_memory=True) return None def sample_async( @@ -656,6 +667,8 @@ def sample_async( new_tokens = self.store.new_tokens finish_reasons = self.store.finish_reasons log_probs_host = self.log_probs_host(scheduled_requests) + top_log_probs_host = self.store.top_log_probs_host + top_tokens_host = self.store.top_tokens_host seq_slots_host = torch.tensor( [r.py_seq_slot for r in requests], dtype=torch.int64, # for index_fill_ @@ -667,7 +680,9 @@ def sample_async( num_context_logits_prefix_sum, seq_slots=seq_slots, seq_slots_host=seq_slots_host, - log_probs_host=log_probs_host) + log_probs_host=log_probs_host, + top_log_probs_host=top_log_probs_host, + top_tokens_host=top_tokens_host) self._write_finish_reasons(requests, finish_reasons=finish_reasons, seq_slots=seq_slots, @@ -675,15 +690,7 @@ def sample_async( new_tokens_host = new_tokens.to(device="cpu", non_blocking=True) finish_reasons_host = finish_reasons.to(device="cpu", non_blocking=True) - # Setup these optional values if possible - log_probs = None - top_log_probs = None - top_tokens = None - if log_probs_host is not None: - log_probs = log_probs_host[..., 0] - if self.calculate_top_logprobs: - top_log_probs = log_probs_host[..., 1:] - top_tokens = new_tokens[..., 1:] + sampler_event = torch.cuda.Event() sampler_event.record() return SampleStateTorch( @@ -693,6 +700,8 @@ def sample_async( new_tokens=new_tokens_host, log_probs=log_probs_host, finish_reasons=finish_reasons_host, + top_tokens=top_tokens_host, + top_log_probs=top_log_probs_host, ), sampler_event=sampler_event, ) @@ -950,7 +959,9 @@ def _process_requests(self, *, seq_slots: torch.Tensor, seq_slots_host: torch.Tensor, - log_probs_host: torch.Tensor | None = None): + log_probs_host: torch.Tensor | None = None, + top_log_probs_host: torch.Tensor | None = None, + top_tokens_host: torch.Tensor | None = None): # raw_logits should contain only the logits from the gen requests. # If return context logits is requested, fetch only the logits from gen requests. @@ -981,8 +992,8 @@ def _process_requests(self, next_tokens = torch.argmax(logits, dim=-1) self.append_eagle3(next_tokens, model_outputs) int_next_tokens = next_tokens.to(torch.int, non_blocking=True) - next_tokens = int_next_tokens.view(1, -1, SINGLE_BEAM_WIDTH, 1) - new_tokens[:1, ..., :1].index_copy_(1, seq_slots, next_tokens) + next_tokens = int_next_tokens.view(1, -1, SINGLE_BEAM_WIDTH) + new_tokens[:1].index_copy_(1, seq_slots, next_tokens) return strategies = sampling_strategies(requests) @@ -1003,7 +1014,7 @@ def _process_requests(self, ] logits = self._apply_embedding_bias(logits, requests, steps_per_request) - batched_next_tokens, batched_softmax, batched_indices = sample( + batched_next_tokens, batched_softmax, batched_all_token_ids = sample( batched_strategy, logits, generator) self.append_eagle3(batched_next_tokens, model_outputs) @@ -1011,8 +1022,9 @@ def _process_requests(self, for i, (strategy, slot, steps, request) in enumerate( zip(strategies, seq_slots_host, num_steps, requests)): # if request.calculate_top_logprobs, then request.py_top_logprobs is the number of top tokens to sample - # otherwise, it is None and should be defaulted to0 - num_top_tokens = request.py_top_logprobs if request.calculate_top_logprobs else 0 + # otherwise, it is None and should be defaulted to 0 + need_top_logprobs = request.calculate_top_logprobs + num_top_tokens = request.py_top_logprobs if need_top_logprobs else 0 if num_top_tokens > self.max_top_logprobs: raise ValueError( f"top_logprobs {request.py_top_logprobs} cannot exceed max_top_logprobs {self.max_top_logprobs}" @@ -1020,40 +1032,44 @@ def _process_requests(self, input_slice = slice(offset, offset + steps) logits = raw_logits[input_slice] - req = requests[i] - if batched_next_tokens is None: - logits = self._apply_embedding_bias(logits, [req]) - next_tokens, softmax, indices = sample(strategy, logits, - generator) + logits = self._apply_embedding_bias(logits, [request]) + next_tokens, softmax, all_token_ids = sample( + strategy, logits, generator) else: # Batched processing already applied bias, just use the results next_tokens = batched_next_tokens[input_slice] softmax = batched_softmax[input_slice] - indices = batched_indices[ - input_slice] if strategy != GREEDY else None - - current_slice = slice(0, steps), slot, BEAM_0, slice( - 0, 1 + num_top_tokens) - if request.calculate_top_logprobs and strategy != GREEDY: - # concatenate the sampled tokens with the top tokens - next_tokens = torch.cat( - [next_tokens.unsqueeze(1), indices[:, :num_top_tokens]], - dim=1) - else: - next_tokens = next_tokens.unsqueeze(1) + all_token_ids = None if not need_top_logprobs else batched_all_token_ids[ + input_slice] + + current_slice = slice(0, steps), slot, BEAM_0 new_tokens[current_slice] = next_tokens if request.py_draft_logits is not None: request.py_target_probs = softmax.clone() if log_probs_host is not None: assert BEAM_0 == 0, "The following call relies on beam_width to be 1 - hence the unsqueeze" - token_probs = torch.gather(softmax, dim=1, - index=next_tokens).reshape( - steps, -1) + token_probs = torch.gather( + softmax, dim=1, + index=next_tokens.unsqueeze(1)).reshape(steps) log_probs = torch.log(token_probs) - log_probs_host[slot, BEAM_0, :steps, :1 + num_top_tokens].copy_( - log_probs, non_blocking=True) + log_probs_host[slot, BEAM_0, :steps].copy_(log_probs, + non_blocking=True) + if need_top_logprobs: + top_token_ids = all_token_ids[:, :num_top_tokens] + top_token_probs = torch.gather(softmax, + dim=1, + index=top_token_ids).reshape( + steps, -1) + top_log_probs = torch.log(top_token_probs) + top_log_probs_host[slot, + BEAM_0, :steps, :num_top_tokens].copy_( + top_log_probs, non_blocking=True) + top_tokens_host[:steps, slot, + BEAM_0, :num_top_tokens].copy_( + top_token_ids, non_blocking=True) + offset += steps diff --git a/tensorrt_llm/_torch/speculative/mtp.py b/tensorrt_llm/_torch/speculative/mtp.py index 733637984827..8d947c0be50b 100644 --- a/tensorrt_llm/_torch/speculative/mtp.py +++ b/tensorrt_llm/_torch/speculative/mtp.py @@ -215,10 +215,11 @@ def prepare(self): class MTPStore(TorchStore): def __init__(self, *, max_draft_len: int, max_num_sequences: int, - max_beam_width: int): + max_beam_width: int, max_top_logprobs: int): super().__init__(max_draft_len=max_draft_len, max_num_sequences=max_num_sequences, - max_beam_width=max_beam_width) + max_beam_width=max_beam_width, + max_top_logprobs=max_top_logprobs) self.next_new_tokens = int_tensor( (self.max_tokens, self.max_num_sequences, SINGLE_BEAM_WIDTH)) self.next_draft_tokens = int_tensor( @@ -241,7 +242,8 @@ def __init__(self, args: TorchSampler.Args, *, nextn: int): self.draft_len = nextn self.store = MTPStore(max_draft_len=nextn, max_num_sequences=args.max_num_sequences, - max_beam_width=args.max_beam_width) + max_beam_width=args.max_beam_width, + max_top_logprobs=args.max_top_logprobs) self.max_seq_len = args.max_seq_len def _request_common_handling(self, request: LlmRequest, diff --git a/tests/unittest/_torch/sampler/test_return_top_k_logprobs.py b/tests/unittest/_torch/sampler/test_return_top_k_logprobs.py index 31a4ba9e9a0f..6533b705333a 100644 --- a/tests/unittest/_torch/sampler/test_return_top_k_logprobs.py +++ b/tests/unittest/_torch/sampler/test_return_top_k_logprobs.py @@ -64,6 +64,15 @@ def test_generate_with_top_logprobs(top_logprobs: int, top_k: int, top_p: float, assert len(output.outputs[0].top_logprobs) == max_tokens assert len(output.outputs[0].logprobs[0]) == 1 assert len(output.outputs[0].top_logprobs[0]) == top_logprobs + sampled_token = list(output.outputs[0].logprobs[0].keys())[0] + sampled_token_logprob = output.outputs[0].logprobs[0][ + sampled_token].logprob + if sampled_token in output.outputs[0].top_logprobs[0]: + assert sampled_token_logprob == output.outputs[ + 0].top_logprobs[0][sampled_token].logprob + else: + assert sampled_token_logprob <= list( + output.outputs[0].top_logprobs[0].values())[-1].logprob else: # expected to fail testcases with pytest.raises(ValueError, match="top_logprobs.*"): @@ -77,7 +86,7 @@ def test_generate_with_top_logprobs(top_logprobs: int, top_k: int, top_p: float, @force_ampere # Save H100 resource @pytest.mark.threadleak(enabled=False) -def test_generate_with_top_logprobs_and_disabler_logprobs( +def test_generate_with_top_logprobs_and_disabled_logprobs( llm_torch, input_prompts): max_tokens = 8 top_logprobs = 2 From a0c04121e7d3455c6411e499338f328a09ad85ff Mon Sep 17 00:00:00 2001 From: Stefan Niebler <82932102+stnie@users.noreply.github.com> Date: Wed, 10 Sep 2025 12:39:06 +0000 Subject: [PATCH 07/10] [TRTLLM-5670][chore] Change top logprobs API to conform more with current LLM-API - Removed the top_logprobs parameter from SamplingParams and related classes to streamline log probability management. - Change _get_output_config to explicitly pass logprobs as top-logprobs to OutputConfig object. - Updated setup_llm function to conditionally set logprobs based on input arguments. - Adjusted CompletionOutput and related tests to reflect the removal of top_logprobs. - Enhanced validation checks to ensure correct usage of logprobs in sampling scenarios. Signed-off-by: Stefan Niebler <82932102+stnie@users.noreply.github.com> --- examples/llm-api/quickstart_advanced.py | 13 ++-- tensorrt_llm/_torch/pyexecutor/_util.py | 5 +- tensorrt_llm/executor/result.py | 10 +-- tensorrt_llm/llmapi/llm.py | 4 -- tensorrt_llm/sampling_params.py | 24 +++---- .../sampler/test_return_top_k_logprobs.py | 66 +++++-------------- .../references/completion_output.yaml | 3 - .../references/sampling_params.yaml | 3 - 8 files changed, 36 insertions(+), 92 deletions(-) diff --git a/examples/llm-api/quickstart_advanced.py b/examples/llm-api/quickstart_advanced.py index 40fb065f644d..9b42c92f5f1c 100644 --- a/examples/llm-api/quickstart_advanced.py +++ b/examples/llm-api/quickstart_advanced.py @@ -223,7 +223,6 @@ def setup_llm(args, **kwargs): 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', @@ -270,6 +269,11 @@ 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, @@ -277,8 +281,7 @@ def setup_llm(args, **kwargs): top_p=args.top_p, return_context_logits=args.return_context_logits, return_generation_logits=args.return_generation_logits, - logprobs=args.logprobs, - top_logprobs=args.top_logprobs, + logprobs=num_logprobs, n=args.n, best_of=best_of, use_beam_search=use_beam_search) @@ -320,10 +323,6 @@ def main(): ) if args.logprobs: print(f"[{i}]{sequence_id_text} Logprobs: {sequence.logprobs}") - if args.top_logprobs: - print( - f"[{i}]{sequence_id_text} Top logprobs: {sequence.top_logprobs}" - ) if __name__ == '__main__': diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index ea558f36b614..a7dfe31f7bf6 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -686,8 +686,8 @@ def create_py_executor_instance( def create_torch_sampler_args(mapping: Mapping, *, max_seq_len: int, - enable_mixed_sampler: bool, - max_top_logprobs: int, 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 @@ -715,7 +715,6 @@ def instantiate_sampler(engine: PyTorchModelEngine, 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) diff --git a/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py index 6f3eadcf614e..21c17cd0e591 100644 --- a/tensorrt_llm/executor/result.py +++ b/tensorrt_llm/executor/result.py @@ -93,7 +93,6 @@ class CompletionOutput: cumulative_logprob (float, optional): The cumulative log probability of the generated output text. Defaults to None. logprobs (TokenLogprobs | List[float], optional): The log probabilities of the top probability words at each position if the logprobs are requested. Defaults to None. prompt_logprobs (TokenLogprobs, optional): The log probabilities per prompt token. Defaults to None. - top_logprobs (TokenLogprobs, optional): The log probabilities of the top probability words at each position if the top_logprobs are requested. Defaults to None. 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. @@ -113,7 +112,6 @@ class CompletionOutput: logprobs: Optional[TokenLogprobs | List[float]] = field(default_factory=list) prompt_logprobs: Optional[TokenLogprobs] = field(default_factory=list) - top_logprobs: Optional[TokenLogprobs] = field(default_factory=list) finish_reason: Optional[Literal['stop', 'length', 'timeout', 'cancelled']] = None stop_reason: Optional[Union[int, str]] = None @@ -263,11 +261,9 @@ def _handle_sequence(self, output.logprobs = output.logprobs[:output.length] assert len(output.logprobs) == output.length if response_tensors.top_log_probs is not None: - output.top_logprobs = response_tensors.top_log_probs[src_idx] - if finish_reasons[src_idx] != tllm.FinishReason.CANCELLED: - if len(output.top_logprobs) > output.length: - output.top_logprobs = output.top_logprobs[:output. - length] + for logprob_idx in range(len(output.logprobs)): + output.logprobs[logprob_idx].update( + response_tensors.top_log_probs[src_idx][logprob_idx]) if response_tensors.generation_logits is not None: output.generation_logits = response_tensors.generation_logits[ src_idx, :output.length] diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py index b2665b587ec8..192181bb1c50 100644 --- a/tensorrt_llm/llmapi/llm.py +++ b/tensorrt_llm/llmapi/llm.py @@ -580,10 +580,6 @@ def _check_arguments(self, prompt_len: int, query_len: int, raise ValueError( f"`prompt_logprobs` in sampling_params is not supported in the PyTorch backend yet. Received `prompt_logprobs={sampling_params.prompt_logprobs}`. Please unset this field." ) - if sampling_params.logprobs and sampling_params.logprobs > 1: - raise ValueError( - f"PyTorch backend currently only supports `logprobs=1`. Received `logprobs={sampling_params.logprobs}` (Top{sampling_params.logprobs} logprobs). Please set `logprobs=1` in `sampling_params` instead." - ) # Check prompt length and query length against max_num_tokens to filter illegal requests. # Skip check for gen-only requests if self.args.backend == "pytorch" and not self.args.enable_chunked_prefill and not is_gen_only: diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index ac635bd63c75..d1cfafb85fb9 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -173,7 +173,6 @@ class SamplingParams: logprobs (int, optional): Number of log probabilities to return per output token. Defaults to None. prompt_logprobs (int, optional): Number of log probabilities to return per prompt token. Defaults to None. - top_logprobs (int, optional): Number of top log probabilities to return per output token. Defaults to None. return_context_logits (bool): Controls if Result should contain the context logits. Defaults to False. return_generation_logits (bool): Controls if Result should contain the generation logits. Defaults to False. exclude_input_from_output (bool): Controls if output tokens in Result should include the input tokens. Defaults to True. @@ -242,7 +241,6 @@ class SamplingParams: # Keep the below fields in sync with tllme.OutputConfig logprobs: Optional[int] = None prompt_logprobs: Optional[int] = None - top_logprobs: Optional[int] = None return_context_logits: bool = False return_generation_logits: bool = False exclude_input_from_output: bool = True @@ -325,25 +323,18 @@ def _validate(self): self.logprobs = self.logprobs and int(self.logprobs) self.prompt_logprobs = self.prompt_logprobs and int(self.prompt_logprobs) - if self.top_logprobs is not None: - if not self.logprobs: - raise ValueError("You need to set logprobs to True to use top_logprobs.") - if self.top_logprobs < 1: - raise ValueError(f"top_logprobs must be >= 1, got {self.top_logprobs}") - - if self._greedy_decoding: + if self.logprobs is not None: + if self._greedy_decoding and self.logprobs > 0: # check for greedy sampling raise ValueError( - "top_logprobs must not be specified for greedy sampling, " - "as it provides the same output as logprobs in this case." + "logprobs > 0 must not be specified for greedy sampling, " + "as it provides the same output as logprobs = 0 in this case." ) - elif self.top_k is not None and self.top_k < self.top_logprobs: + elif self.top_k is not None and self.top_k < self.logprobs: # check for top-k sampling - raise ValueError( - f"top_logprobs {self.top_logprobs} must not exceed top_k {self.top_k}" - ) + raise ValueError(f"logprobs {self.logprobs} must not exceed top_k {self.top_k}") if self.use_beam_search: - raise ValueError("top_logprobs + beam search is not supported") + raise ValueError("logprobs > 0 + beam search is not supported") @property def _greedy_decoding(self) -> bool: @@ -471,6 +462,7 @@ def _get_output_config(self, is_pytorch_backend: bool = False) -> tllme.OutputCo if is_pytorch_backend: config_kwargs["return_log_probs"] = bool(self.logprobs) + config_kwargs["top_logprobs"] = self.logprobs else: config_kwargs["return_log_probs"] = self._return_log_probs diff --git a/tests/unittest/_torch/sampler/test_return_top_k_logprobs.py b/tests/unittest/_torch/sampler/test_return_top_k_logprobs.py index 6533b705333a..6cbac43280c2 100644 --- a/tests/unittest/_torch/sampler/test_return_top_k_logprobs.py +++ b/tests/unittest/_torch/sampler/test_return_top_k_logprobs.py @@ -36,7 +36,7 @@ def llm_torch(input_prompts): @force_ampere # Save H100 resource -@pytest.mark.parametrize("top_logprobs", [None, 0, 2, 4]) +@pytest.mark.parametrize("top_logprobs", [0, 2, 4]) @pytest.mark.parametrize("top_k", [None, 2]) @pytest.mark.parametrize("top_p", [None, 0.5]) @pytest.mark.threadleak(enabled=False) @@ -45,61 +45,29 @@ def test_generate_with_top_logprobs(top_logprobs: int, top_k: int, top_p: float, max_tokens = 8 is_topk = top_k is not None and top_k > 1 is_topp = top_p is not None and top_p > 0.0 and not is_topk - uses_top_logprobs = top_logprobs is not None - is_valid_top_logprobs = top_logprobs > 0 if uses_top_logprobs else True - is_valid_setup = not uses_top_logprobs or (is_valid_top_logprobs and ( - (is_topk and top_logprobs <= top_k) or is_topp)) + is_valid_setup = (is_topk + and top_logprobs <= top_k) or is_topp or top_logprobs == 0 if is_valid_setup: # passing testcases sampling_params = SamplingParams(max_tokens=max_tokens, - logprobs=True, + logprobs=top_logprobs, top_k=top_k, top_p=top_p, - temperature=1.0, - top_logprobs=top_logprobs) + temperature=1.0) for output in llm_torch.generate(input_prompts, sampling_params=sampling_params): - if top_logprobs is not None and top_logprobs > 0: + if top_logprobs > 0: assert len(output.outputs[0].logprobs) == max_tokens - assert len(output.outputs[0].top_logprobs) == max_tokens - assert len(output.outputs[0].logprobs[0]) == 1 - assert len(output.outputs[0].top_logprobs[0]) == top_logprobs - sampled_token = list(output.outputs[0].logprobs[0].keys())[0] - sampled_token_logprob = output.outputs[0].logprobs[0][ - sampled_token].logprob - if sampled_token in output.outputs[0].top_logprobs[0]: - assert sampled_token_logprob == output.outputs[ - 0].top_logprobs[0][sampled_token].logprob - else: - assert sampled_token_logprob <= list( - output.outputs[0].top_logprobs[0].values())[-1].logprob + assert len( + output.outputs[0].logprobs[0]) == top_logprobs or len( + output.outputs[0].logprobs[0]) == top_logprobs + 1 else: # expected to fail testcases - with pytest.raises(ValueError, match="top_logprobs.*"): - sampling_params = SamplingParams(max_tokens=max_tokens, - logprobs=True, - top_k=top_k, - top_p=top_p, - temperature=1.0, - top_logprobs=top_logprobs) - - -@force_ampere # Save H100 resource -@pytest.mark.threadleak(enabled=False) -def test_generate_with_top_logprobs_and_disabled_logprobs( - llm_torch, input_prompts): - max_tokens = 8 - top_logprobs = 2 - top_k = 2 - top_p = None - - # expected to fail testcases - with pytest.raises( - ValueError, - match=".*You need to set logprobs to True to use top_logprobs.*"): - sampling_params = SamplingParams(max_tokens=max_tokens, - logprobs=False, - top_k=top_k, - top_p=top_p, - temperature=1.0, - top_logprobs=top_logprobs) + with pytest.raises(ValueError, match="logprobs.*"): + sampling_params = SamplingParams( + max_tokens=max_tokens, + logprobs=top_logprobs, + top_k=top_k, + top_p=top_p, + temperature=1.0, + ) diff --git a/tests/unittest/api_stability/references/completion_output.yaml b/tests/unittest/api_stability/references/completion_output.yaml index 06159b2446eb..cd1f49a89cf4 100644 --- a/tests/unittest/api_stability/references/completion_output.yaml +++ b/tests/unittest/api_stability/references/completion_output.yaml @@ -10,9 +10,6 @@ methods: request_perf_metrics: annotation: Optional[tensorrt_llm.bindings.executor.RequestPerfMetrics] default: null - top_logprobs: - annotation: Optional[list[dict[int, tensorrt_llm.executor.result.Logprob]]] - default: null return_annotation: None properties: length: diff --git a/tests/unittest/api_stability/references/sampling_params.yaml b/tests/unittest/api_stability/references/sampling_params.yaml index f460f264a0e9..1ac178ed755a 100644 --- a/tests/unittest/api_stability/references/sampling_params.yaml +++ b/tests/unittest/api_stability/references/sampling_params.yaml @@ -11,8 +11,5 @@ methods: beam_width_array: annotation: Optional[List[int]] default: null - top_logprobs: - annotation: Optional[int] - default: null return_annotation: None properties: {} From 6d31b22a9901c62f9a7a6f520d17094496ad5f37 Mon Sep 17 00:00:00 2001 From: Stefan Niebler <82932102+stnie@users.noreply.github.com> Date: Mon, 15 Sep 2025 08:10:17 +0000 Subject: [PATCH 08/10] [TRTLLM-5670][chore] Update logprobs validation for greedy sampling - Adjusted the validation logic in SamplingParams to ensure that logprobs must not exceed 1 when greedy sampling is enabled. - Updated error message for clarity regarding the constraints on logprobs in greedy sampling scenarios. Signed-off-by: Stefan Niebler <82932102+stnie@users.noreply.github.com> --- tensorrt_llm/sampling_params.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index d1cfafb85fb9..1e3aec17c611 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -324,12 +324,9 @@ def _validate(self): self.prompt_logprobs = self.prompt_logprobs and int(self.prompt_logprobs) if self.logprobs is not None: - if self._greedy_decoding and self.logprobs > 0: + if self._greedy_decoding and self.logprobs > 1: # check for greedy sampling - raise ValueError( - "logprobs > 0 must not be specified for greedy sampling, " - "as it provides the same output as logprobs = 0 in this case." - ) + raise ValueError(f"logprobs {self.logprobs} must not exceed 1 for greedy sampling") elif self.top_k is not None and self.top_k < self.logprobs: # check for top-k sampling raise ValueError(f"logprobs {self.logprobs} must not exceed top_k {self.top_k}") From 98b472a14083adb7c01e6cb65079dd3fab6a1b2e Mon Sep 17 00:00:00 2001 From: Stefan Niebler <82932102+stnie@users.noreply.github.com> Date: Mon, 15 Sep 2025 15:38:22 +0000 Subject: [PATCH 09/10] [TRTLLM-5670][chore] Update Greedy sampling to also support top logprobs of 1 - Modified the return statement in `greedy_search_sampling_batch` to include `next_tokens` as an unsqueezed tensor, ensuring compatibility with downstream processing. - Added `max_top_logprobs` parameter to the configuration in the test file to ensure correct execution. Signed-off-by: Stefan Niebler <82932102+stnie@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/sampler.py | 2 +- tests/unittest/llmapi/apps/_test_trtllm_serve_top_logprobs.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py index b52e1d24080c..69754e2ece54 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler.py @@ -292,7 +292,7 @@ def greedy_search_sampling_batch(logits): """ next_tokens = torch.argmax(logits, dim=-1) softmax = torch.softmax(logits, dim=-1) - return next_tokens, softmax, None + return next_tokens, softmax, next_tokens.unsqueeze(-1) def get_rejected_indices(draft_probs: torch.Tensor, target_probs: torch.Tensor, diff --git a/tests/unittest/llmapi/apps/_test_trtllm_serve_top_logprobs.py b/tests/unittest/llmapi/apps/_test_trtllm_serve_top_logprobs.py index d287e5e35eb5..209ca8840505 100644 --- a/tests/unittest/llmapi/apps/_test_trtllm_serve_top_logprobs.py +++ b/tests/unittest/llmapi/apps/_test_trtllm_serve_top_logprobs.py @@ -31,7 +31,8 @@ def temp_extra_llm_api_options_file(): "gather_generation_logits": True, "kv_cache_config": { "enable_block_reuse": False, - } + }, + "max_top_logprobs": 1 } with open(temp_file_path, 'w') as f: From 31bc186aa8f5485664840d78316728a68552d279 Mon Sep 17 00:00:00 2001 From: Stefan Niebler <82932102+stnie@users.noreply.github.com> Date: Wed, 17 Sep 2025 08:27:50 +0000 Subject: [PATCH 10/10] [TRTLLM-5670][chore] Adjust top logprobs handling in tests - Updated test_embedding_bias_with_torch_sampler_strategies cases to correctly handle top logprobs of 1 in the parameterization - Enhanced validation logic in test_return_top_k_logprobs to accommodate the new conditions for valid setups in sampling tests (greedy sampling). Signed-off-by: Stefan Niebler <82932102+stnie@users.noreply.github.com> --- .../unittest/_torch/sampler/test_return_top_k_logprobs.py | 8 +++++--- tests/unittest/llmapi/test_llm_pytorch.py | 5 ++++- 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/tests/unittest/_torch/sampler/test_return_top_k_logprobs.py b/tests/unittest/_torch/sampler/test_return_top_k_logprobs.py index 6cbac43280c2..6b07bf65f453 100644 --- a/tests/unittest/_torch/sampler/test_return_top_k_logprobs.py +++ b/tests/unittest/_torch/sampler/test_return_top_k_logprobs.py @@ -36,7 +36,7 @@ def llm_torch(input_prompts): @force_ampere # Save H100 resource -@pytest.mark.parametrize("top_logprobs", [0, 2, 4]) +@pytest.mark.parametrize("top_logprobs", [0, 1, 4]) @pytest.mark.parametrize("top_k", [None, 2]) @pytest.mark.parametrize("top_p", [None, 0.5]) @pytest.mark.threadleak(enabled=False) @@ -45,8 +45,10 @@ def test_generate_with_top_logprobs(top_logprobs: int, top_k: int, top_p: float, max_tokens = 8 is_topk = top_k is not None and top_k > 1 is_topp = top_p is not None and top_p > 0.0 and not is_topk - is_valid_setup = (is_topk - and top_logprobs <= top_k) or is_topp or top_logprobs == 0 + is_greedy = not (is_topk or is_topp) + is_valid_setup = (is_topk and top_logprobs + <= top_k) or is_topp or top_logprobs == 0 or ( + is_greedy and top_logprobs == 1) if is_valid_setup: # passing testcases sampling_params = SamplingParams(max_tokens=max_tokens, diff --git a/tests/unittest/llmapi/test_llm_pytorch.py b/tests/unittest/llmapi/test_llm_pytorch.py index e52d2cae6c69..7381b5c04a78 100644 --- a/tests/unittest/llmapi/test_llm_pytorch.py +++ b/tests/unittest/llmapi/test_llm_pytorch.py @@ -244,6 +244,7 @@ def test_embedding_bias_with_torch_sampler_strategies(enable_mixed_sampler, vocab_size_padded = 32000 embedding_bias = torch.zeros(vocab_size_padded) embedding_bias[biased_word_id] = torch.finfo(torch.float32).max + max_top_logprobs = 0 sampling_kwargs = { "max_tokens": 6, @@ -252,6 +253,7 @@ def test_embedding_bias_with_torch_sampler_strategies(enable_mixed_sampler, if enable_logprobs: sampling_kwargs["logprobs"] = 1 + max_top_logprobs = 1 # All test cases use greedy sampling for simplicity sampling_params = SamplingParams(**sampling_kwargs) @@ -260,7 +262,8 @@ def test_embedding_bias_with_torch_sampler_strategies(enable_mixed_sampler, prompts, ["Z Z Z Z Z Z"], sampling_params=sampling_params, backend="pytorch", - enable_mixed_sampler=enable_mixed_sampler) + enable_mixed_sampler=enable_mixed_sampler, + max_top_logprobs=max_top_logprobs) def llama_7b_lora_from_dir_test_harness(**llm_kwargs) -> None: