tensorrt_llm/_torch/auto_deploy/shim/demollm.py (1)

16-17: Sampling defaults ignore top_p/temperature paths

With the new semantics, any request that only tweaks temperature or top_p (leaving top_k unset/≤1) should still perform stochastic sampling. This branch forces those cases down the greedy path, so temperature-only or top-p-only configs silently stop sampling. Please route through the full strategy logic (greedy vs temperature vs top-p vs top-k vs top-k+top-p) instead of keying solely on top_k > 1.

Apply this diff to align the shim with the dispatcher logic:

-from ....sampling_params import SamplingParams
-from ...pyexecutor.sampler import greedy_search_sampling_batch, top_k_sampling_batch
+from ....sampling_params import SamplingParams
+from ...pyexecutor.sampler import (
+    greedy_search_sampling_batch,
+    temperature_sampling_batch,
+    top_k_sampling_batch,
+    top_k_top_p_sampling_batch,
+    top_p_sampling_batch,
+)
@@
-        if isinstance(sampling_params.top_k, int) and sampling_params.top_k > 1:
-            idx_next, probs = top_k_sampling_batch(
-                logits, top_k=sampling_params.top_k, temperature=1.0
-            )
-        else:
-            idx_next, probs = greedy_search_sampling_batch(logits)
+        if SamplingParams.params_imply_greedy_decoding(
+            temperature=sampling_params.temperature,
+            top_p=sampling_params.top_p,
+            top_k=sampling_params.top_k,
+        ):
+            idx_next, probs = greedy_search_sampling_batch(logits)
+        else:
+            vocab_size = logits.size(-1)
+            temperature = sampling_params.temperature or 1.0
+            top_p = sampling_params.top_p or 1.0
+            top_k = sampling_params.top_k or vocab_size
+
+            if top_p < 1 and top_k < vocab_size:
+                idx_next, probs = top_k_top_p_sampling_batch(
+                    logits, top_k=top_k, top_p=top_p, temperature=temperature
+                )
+            elif top_p < 1:
+                idx_next, probs = top_p_sampling_batch(
+                    logits, top_p=top_p, temperature=temperature
+                )
+            elif top_k < vocab_size:
+                idx_next, probs = top_k_sampling_batch(
+                    logits, top_k=top_k, temperature=temperature
+                )
+            else:
+                idx_next, probs = temperature_sampling_batch(
+                    logits, temperature=temperature
+                )

Also applies to: 237-244

tensorrt_llm/_torch/pyexecutor/sampler.py (2)

502-504: Advanced indexing requires Long indices.

softmax[...] with an index tensor needs dtype=torch.long. Casting to device alone isn’t enough when current dtype is int32.

-    if filter_softmax and softmax_indices is not None:
-        softmax = softmax[softmax_indices.to(softmax.device, non_blocking=True)]
+    if filter_softmax and softmax_indices is not None:
+        idx = softmax_indices.to(dtype=torch.long,
+                                 device=softmax.device,
+                                 non_blocking=True)
+        softmax = softmax[idx]

---

1-1: Add NVIDIA Apache-2.0 license header

The file tensorrt_llm/_torch/pyexecutor/sampler.py is missing the required NVIDIA Apache-2.0 header; prepend it as per CODING_GUIDELINES.md:

+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.