diff --git a/cpp/include/tensorrt_llm/executor/executor.h b/cpp/include/tensorrt_llm/executor/executor.h
index 9dda07d19c6a..6bc0b5f247d5 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<std::vector<AdditionalModelOutput>> additionalModelOutputs = std::nullopt);
+        std::optional<std::vector<AdditionalModelOutput>> additionalModelOutputs = std::nullopt,
+        std::optional<SizeType32> topLogProbs = std::nullopt);
 
     /// @brief Controls if Result should contain log probabilities. Default is false.
     bool returnLogProbs;
@@ -233,6 +234,8 @@ class OutputConfig
 
     /// @brief The additional outputs to gather from the model.
     std::optional<std::vector<AdditionalModelOutput>> additionalModelOutputs;
+    /// @brief The number of logprobs to return. Default is 0.
+    std::optional<SizeType32> topLogProbs;
 };
 
 /// @brief Configuration for speculative decoding with external draft tokens.
diff --git a/cpp/tensorrt_llm/executor/outputConfig.cpp b/cpp/tensorrt_llm/executor/outputConfig.cpp
index 43fad705cab9..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<std::vector<AdditionalModelOutput>> additionalModelOutputs)
+    std::optional<std::vector<AdditionalModelOutput>> additionalModelOutputs, std::optional<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..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<bool>(state[0]), nb::cast<bool>(state[1]),
             nb::cast<bool>(state[2]), nb::cast<bool>(state[3]), nb::cast<bool>(state[4]), nb::cast<bool>(state[5]),
-            nb::cast<std::optional<std::vector<tle::AdditionalModelOutput>>>(state[6]));
+            nb::cast<std::optional<std::vector<tle::AdditionalModelOutput>>>(state[6]),
+            nb::cast<std::optional<SizeType32>>(state[7]));
     };
     nb::class_<tle::OutputConfig>(m, "OutputConfig")
         .def(
@@ -216,17 +218,18 @@ void initRequestBindings(nb::module_& m)
             [](tle::OutputConfig& self, std::optional<bool> return_log_probs, std::optional<bool> return_context_logits,
                 std::optional<bool> return_generation_logits, std::optional<bool> exclude_input_from_output,
                 std::optional<bool> return_encoder_output, std::optional<bool> return_perf_metrics,
-                std::optional<std::vector<tle::AdditionalModelOutput>> additional_model_outputs)
+                std::optional<std::vector<tle::AdditionalModelOutput>> additional_model_outputs,
+                std::optional<SizeType32> top_logprobs)
             {
                 new (&self) tle::OutputConfig(return_log_probs.value_or(false), return_context_logits.value_or(false),
                     return_generation_logits.value_or(false), exclude_input_from_output.value_or(false),
                     return_encoder_output.value_or(false), return_perf_metrics.value_or(false),
-                    additional_model_outputs);
+                    additional_model_outputs, top_logprobs);
             },
             nb::arg("return_log_probs") = nb::none(), nb::arg("return_context_logits") = nb::none(),
             nb::arg("return_generation_logits") = nb::none(), nb::arg("exclude_input_from_output") = nb::none(),
             nb::arg("return_encoder_output") = nb::none(), nb::arg("return_perf_metrics") = nb::none(),
-            nb::arg("additional_model_outputs") = nb::none())
+            nb::arg("additional_model_outputs") = nb::none(), nb::arg("top_logprobs") = nb::none())
         .def_rw("return_log_probs", &tle::OutputConfig::returnLogProbs)
         .def_rw("return_context_logits", &tle::OutputConfig::returnContextLogits)
         .def_rw("return_generation_logits", &tle::OutputConfig::returnGenerationLogits)
@@ -234,6 +237,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..55ae7fb503a5 100644
--- a/cpp/tensorrt_llm/pybind/executor/request.cpp
+++ b/cpp/tensorrt_llm/pybind/executor/request.cpp
@@ -191,24 +191,27 @@ void initRequestBindings(pybind11::module_& m)
     auto outputConfigGetstate = [](tle::OutputConfig const& self)
     {
         return py::make_tuple(self.returnLogProbs, self.returnContextLogits, self.returnGenerationLogits,
-            self.excludeInputFromOutput, self.returnEncoderOutput, self.returnPerfMetrics, self.additionalModelOutputs);
+            self.excludeInputFromOutput, self.returnEncoderOutput, self.returnPerfMetrics, self.additionalModelOutputs,
+            self.topLogProbs);
     };
     auto outputConfigSetstate = [](py::tuple const& state)
     {
-        if (state.size() != 7)
+        if (state.size() != 8)
         {
             throw std::runtime_error("Invalid OutputConfig state!");
         }
         return tle::OutputConfig(state[0].cast<bool>(), state[1].cast<bool>(), state[2].cast<bool>(),
             state[3].cast<bool>(), state[4].cast<bool>(), state[5].cast<bool>(),
-            state[6].cast<std::optional<std::vector<tle::AdditionalModelOutput>>>());
+            state[6].cast<std::optional<std::vector<tle::AdditionalModelOutput>>>(),
+            state[7].cast<std::optional<SizeType32>>());
     };
     py::class_<tle::OutputConfig>(m, "OutputConfig")
-        .def(py::init<bool, bool, bool, bool, bool, bool, std::optional<std::vector<tle::AdditionalModelOutput>>>(),
+        .def(py::init<bool, bool, bool, bool, bool, bool, std::optional<std::vector<tle::AdditionalModelOutput>>,
+                 std::optional<SizeType32>>(),
             py::arg("return_log_probs") = false, py::arg("return_context_logits") = false,
             py::arg("return_generation_logits") = false, py::arg("exclude_input_from_output") = false,
             py::arg("return_encoder_output") = false, py::arg("return_perf_metrics") = false,
-            py::arg("additional_model_outputs") = py::none())
+            py::arg("additional_model_outputs") = py::none(), py::arg("top_logprobs") = py::none())
         .def_readwrite("return_log_probs", &tle::OutputConfig::returnLogProbs)
         .def_readwrite("return_context_logits", &tle::OutputConfig::returnContextLogits)
         .def_readwrite("return_generation_logits", &tle::OutputConfig::returnGenerationLogits)
@@ -216,6 +219,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..9b42c92f5f1c 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=None)
+    parser.add_argument('--max_top_logprobs', type=int, default=0)
     return parser
 
 
@@ -213,6 +215,14 @@ def setup_llm(args, **kwargs):
         batching_wait_iters=args.attention_dp_batching_wait_iters,
     )
 
+    # if only top_logprobs is set, set max_top_logprobs to top_logprobs
+    if args.top_logprobs is not None and args.max_top_logprobs < args.top_logprobs:
+        args.max_top_logprobs = args.top_logprobs
+    # Remove this once torch sampler stops using enable_mixed_sampler
+    is_greedy = (not args.max_beam_width > 1) and (
+        args.top_k is None or args.top_k == 1) and (args.top_p is None
+                                                    or args.top_p == 0.0)
+
     llm = LLM(
         model=args.model_dir,
         backend='pytorch',
@@ -246,6 +256,8 @@ def setup_llm(args, **kwargs):
         trust_remote_code=args.trust_remote_code,
         gather_generation_logits=args.return_generation_logits,
         max_beam_width=args.max_beam_width,
+        max_top_logprobs=args.max_top_logprobs,
+        enable_mixed_sampler=not is_greedy,
         **kwargs)
 
     use_beam_search = args.max_beam_width > 1
@@ -257,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,
@@ -264,7 +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,
+        logprobs=num_logprobs,
         n=args.n,
         best_of=best_of,
         use_beam_search=use_beam_search)
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
 
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 684158a80641..a7dfe31f7bf6 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,7 @@ def instantiate_sampler(engine: PyTorchModelEngine,
         mapping,
         max_seq_len=engine.max_seq_len,
         enable_mixed_sampler=pytorch_backend_config.enable_mixed_sampler,
+        max_top_logprobs=pytorch_backend_config.max_top_logprobs,
         max_batch_size=max_batch_size,
         speculative_config=speculative_config,
         max_beam_width=max_beam_width)
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..6078322229d5 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],
@@ -302,6 +326,7 @@ def __init__(
             is_draft: bool = False,
             seq_slot: Optional[int] = None,
             target_seq_slot: Optional[int] = None,
+            top_logprobs: Optional[int] = None,
             **kwargs):
 
         self.py_logits_post_processors = kwargs.pop("py_logits_post_processors",
@@ -346,6 +371,7 @@ def __init__(
             TaskLayerModuleConfig] | None = None
 
         self.py_return_log_probs = return_log_probs
+        self.py_top_logprobs = top_logprobs
         self.py_return_context_logits = return_context_logits
         self.py_return_generation_logits = return_generation_logits
         self.py_return_logits_device_memory = return_logits_device_memory
@@ -394,6 +420,10 @@ def create_response(
     def is_dummy(self):
         return self.is_attention_dp_dummy or self.is_cuda_graph_dummy or self.is_dummy_request
 
+    @property
+    def calculate_top_logprobs(self):
+        return self.py_top_logprobs is not None and self.py_top_logprobs > 0
+
     def finish_by(self, reason: FinishReason, beam: int) -> None:
         """CPP finish by reason does not support beam_width > 1"""
         self.state = LlmRequestState.GENERATION_COMPLETE
@@ -564,7 +594,8 @@ def executor_request_to_llm_request(
         cache_salt_id=executor_request.cache_salt_id,
         arrival_time=getattr(executor_request, "py_arrival_time", None),
         py_multimodal_data=getattr(executor_request, "py_multimodal_data",
-                                   None))
+                                   None),
+        top_logprobs=executor_request.output_config.top_logprobs)
     if child_req_ids:
         for child_id in child_req_ids:
             llm_request.create_child_request(child_id)
diff --git a/tensorrt_llm/_torch/pyexecutor/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler.py
index 1e2897f9adc9..69754e2ece54 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()
@@ -156,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)
@@ -177,13 +185,19 @@ 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,
                          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)
@@ -214,7 +228,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,
@@ -222,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 after applying top-k then top-p mask)
+    """
     logits_dim = logits.dim()
     if logits_dim == 1:
         logits = logits.unsqueeze(0)
@@ -260,13 +280,19 @@ 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):
+    """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
+    return next_tokens, softmax, next_tokens.unsqueeze(-1)
 
 
 def get_rejected_indices(draft_probs: torch.Tensor, target_probs: torch.Tensor,
@@ -362,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))
@@ -390,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)
@@ -414,11 +454,13 @@ 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):
         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.
@@ -426,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
 
@@ -449,19 +492,42 @@ 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, *,
+    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.calculate_top_logprobs:
+            assert state.host.log_probs is not None
+            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=1)
+                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]
+            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])
+            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)`"""
@@ -582,8 +648,9 @@ 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: In lockstep with TRTLLMSampler: https://github.com/NVIDIA/TensorRT-LLM/blob/cea5dd1e3883b18bf50901a7f196f50a9544c28c/cpp/include/tensorrt_llm/runtime/decoderState.h#L103"""
+    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(
@@ -600,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_
@@ -611,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,
@@ -619,6 +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)
+
         sampler_event = torch.cuda.Event()
         sampler_event.record()
         return SampleStateTorch(
@@ -628,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,
         )
@@ -885,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.
@@ -938,25 +1014,35 @@ 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_all_token_ids = 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.calculate_top_logprobs, then request.py_top_logprobs is the number of top tokens to sample
+            # 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}"
+                )
             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 = 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]
+                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:
@@ -964,10 +1050,26 @@ def _process_requests(self,
             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)
+                    softmax, dim=1,
+                    index=next_tokens.unsqueeze(1)).reshape(steps)
+
                 log_probs = torch.log(token_probs)
                 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/tensorrt_llm/executor/result.py b/tensorrt_llm/executor/result.py
index ddbe2f636aaf..21c17cd0e591 100644
--- a/tensorrt_llm/executor/result.py
+++ b/tensorrt_llm/executor/result.py
@@ -260,6 +260,10 @@ 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:
+                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/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..1e3aec17c611 100644
--- a/tensorrt_llm/sampling_params.py
+++ b/tensorrt_llm/sampling_params.py
@@ -323,6 +323,16 @@ def _validate(self):
         self.logprobs = self.logprobs and int(self.logprobs)
         self.prompt_logprobs = self.prompt_logprobs and int(self.prompt_logprobs)
 
+        if self.logprobs is not None:
+            if self._greedy_decoding and self.logprobs > 1:
+                # check for greedy sampling
+                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}")
+            if self.use_beam_search:
+                raise ValueError("logprobs > 0 + beam search is not supported")
+
     @property
     def _greedy_decoding(self) -> bool:
         return (
@@ -449,6 +459,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/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..6b07bf65f453
--- /dev/null
+++ b/tests/unittest/_torch/sampler/test_return_top_k_logprobs.py
@@ -0,0 +1,75 @@
+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, 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_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,
+                                         logprobs=top_logprobs,
+                                         top_k=top_k,
+                                         top_p=top_p,
+                                         temperature=1.0)
+        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].logprobs[0]) == top_logprobs or len(
+                        output.outputs[0].logprobs[0]) == top_logprobs + 1
+    else:
+        # expected to fail testcases
+        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/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/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
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():
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:
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: