[TRTLLM-13212][refactor] Unify sampler ops and clean up dead code in TorchSampler

cpp/tensorrt_llm/thop/dynamicTreeOp.cpp

@@ -25,13 +25,6 @@ namespace tk = tensorrt_llm::kernels::speculative_decoding;
 
 TRTLLM_NAMESPACE_BEGIN
 
-namespace kernels::speculative_decoding
-{
-th::Tensor computeProbsFromLogits(th::Tensor const& logits, th::Tensor const& temperatures,
-    th::optional<th::Tensor> const& topK, th::optional<th::Tensor> const& topP, bool skipTemperature,
-    runtime::SizeType32 kMax);
-} // namespace kernels::speculative_decoding
-
 namespace torch_ext
 {
 
@@ -126,26 +119,6 @@ void verify_dynamic_tree_greedy_out_packed_op(th::Tensor& candidates, th::Tensor
         targetPredict.data_ptr<int32_t>(), treeValid.data_ptr<bool>(), batchSize, numDraftTokens, numSpecStep, stream);
 }
 
-th::Tensor compute_probs_from_logits_op(th::Tensor logits, th::Tensor temperatures, th::optional<th::Tensor> topK,
-    th::optional<th::Tensor> topP, bool skipTemperature)
-{
-    TORCH_CHECK(logits.is_cuda(), "logits must be a CUDA tensor");
-    TORCH_CHECK(temperatures.is_cuda(), "temperatures must be a CUDA tensor");
-    TORCH_CHECK(logits.dim() == 2, "logits must be a 2D tensor");
-    TORCH_CHECK(temperatures.dim() == 1, "temperatures must be a 1D tensor");
-    TORCH_CHECK(logits.size(0) == temperatures.size(0), "logits and temperatures size mismatch");
-    if (topK.has_value() && topK->defined())
-    {
-        TORCH_CHECK(topK->is_cuda(), "top_k must be a CUDA tensor");
-    }
-    if (topP.has_value() && topP->defined())
-    {
-        TORCH_CHECK(topP->is_cuda(), "top_p must be a CUDA tensor");
-    }
-
-    return tk::computeProbsFromLogits(logits, temperatures, topK, topP, skipTemperature, /*kMax=*/0);
-}
-
 //! \brief Target-only rejection sampling verify op (no draft probabilities needed).
 void verify_dynamic_tree_rejection_out_op(th::Tensor& draftTokens, th::Tensor& targetProbs,
     th::Tensor& retrieveNextToken, th::Tensor& retrieveNextSibling, th::Tensor& treeValid, th::Tensor& acceptIndex,
@@ -286,16 +259,3 @@ TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
 {
     m.impl("verify_dynamic_tree_rejection_out_op", &tensorrt_llm::torch_ext::verify_dynamic_tree_rejection_out_op);
 }
-
-TORCH_LIBRARY_FRAGMENT(trtllm, m)
-{
-    m.def(
-        "compute_probs_from_logits_op("
-        "Tensor logits, Tensor temperatures, Tensor? top_k=None, Tensor? top_p=None, "
-        "bool skip_temperature=False) -> Tensor");
-}
-
-TORCH_LIBRARY_IMPL(trtllm, CUDA, m)
-{
-    m.impl("compute_probs_from_logits_op", &tensorrt_llm::torch_ext::compute_probs_from_logits_op);
-}

tensorrt_llm/_torch/auto_deploy/shim/demollm.py

@@ -24,7 +24,7 @@
 
 from tensorrt_llm._torch.pyexecutor.sampler.sampling_utils import (
     greedy_search_sampling_batch,
-    top_k_sampling_batch,
+    top_k_top_p_sampling_batch,
 )
 from tensorrt_llm.executor import GenerationExecutor
 from tensorrt_llm.executor.request import GenerationRequest
@@ -312,7 +312,7 @@ 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) and sampling_params.top_k > 1:
-            idx_next, probs = top_k_sampling_batch(
+            idx_next, probs = top_k_top_p_sampling_batch(
                 logits, top_k=sampling_params.top_k, temperature=1.0
             )
         else:

tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py

@@ -1388,11 +1388,3 @@ def _(like: torch.Tensor,
         out_shape = shape if shape is not None else list(like.shape)
         dtype = out_dtype if out_dtype is not None else like.dtype
         return like.new_empty(out_shape, dtype=dtype), output_buffer_kind
-
-    @torch.library.register_fake("trtllm::compute_probs_from_logits_op")
-    def _(logits: torch.Tensor,
-          temperatures: torch.Tensor,
-          top_k: Optional[torch.Tensor] = None,
-          top_p: Optional[torch.Tensor] = None,
-          skip_temperature: bool = False) -> torch.Tensor:
-        return logits.new_empty(list(logits.shape), dtype=torch.float32)

tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py

@@ -16,9 +16,18 @@
 
 These ops depend on flashinfer; the import is guarded so the module stays
 importable without it.
+
+Randomness can be supplied either way (flashinfer accepts both in one
+signature; explicit ``seed``/``offset`` take precedence over ``generator``):
+
+- ``generator``: stateful host-side ``torch.Generator``, for eager paths.
+- ``seed``/``offset``: stateless device tensors, required under CUDA graph
+  capture (a ``torch.Generator`` advances host-side at launch time, so its
+  state would be frozen into the graph and every replay would reuse the same
+  random values).
 """
 
-from typing import Optional
+from typing import Optional, Union
 
 import torch
 
@@ -27,119 +36,146 @@
 if IS_FLASHINFER_AVAILABLE:
     import flashinfer.sampling
 
+SeedOrTensor = Union[int, torch.Tensor]
+
 
 def top_k_top_p_sampling_from_logits_op(
     logits: torch.Tensor,
     top_k: torch.Tensor,
     top_p: torch.Tensor,
-    seed: Optional[int] = None,
-    offset: Optional[int] = None,
+    *,
+    generator: Optional[torch.Generator] = None,
+    seed: Optional[SeedOrTensor] = None,
+    offset: Optional[SeedOrTensor] = None,
+    check_nan: bool = False,
 ) -> torch.Tensor:
+    """Fused top-k + top-p sampling from pre-softmax logits.
+
+    Randomness: pass ``generator`` (eager) or ``seed``/``offset`` (CUDA graph);
+    see module docstring for the full contract.
+    """
     tokens: torch.Tensor = flashinfer.sampling.top_k_top_p_sampling_from_logits(
-        logits, top_k, top_p, seed=seed, offset=offset
+        logits,
+        top_k=top_k,
+        top_p=top_p,
+        filter_apply_order="top_k_first",
+        deterministic=True,
+        check_nan=check_nan,
+        generator=generator,
+        seed=seed,
+        offset=offset,
     )
     return tokens
 
 
 def sampling_from_probs_op(
     probs: torch.Tensor,
-    seed: Optional[torch.Tensor] = None,
-    offset: Optional[torch.Tensor] = None,
-) -> torch.Tensor:
-    tokens: torch.Tensor = flashinfer.sampling.sampling_from_probs(
-        probs, deterministic=True, seed=seed, offset=offset
-    )
-    return tokens
-
-
-def softmax_op(
-    logits: torch.Tensor,
-    temperature: Optional[torch.Tensor],
-) -> torch.Tensor:
-    probs: torch.Tensor = flashinfer.sampling.softmax(
-        logits, temperature, enable_pdl=get_env_enable_pdl()
-    )
-    return probs
-
-
-def top_k_mask_logits_op(
-    logits: torch.Tensor,
-    top_k: torch.Tensor,
-) -> torch.Tensor:
-    masked: torch.Tensor = flashinfer.sampling.top_k_mask_logits(logits, top_k)
-    return masked
-
-
-def top_p_renorm_probs_op(
-    probs: torch.Tensor,
-    top_p: torch.Tensor,
-) -> torch.Tensor:
-    renormed: torch.Tensor = flashinfer.sampling.top_p_renorm_probs(probs, top_p)
-    return renormed
-
-
-def sampling_from_probs_generator_op(
-    probs: torch.Tensor,
-    generator: Optional[torch.Generator],
+    *,
+    generator: Optional[torch.Generator] = None,
+    seed: Optional[SeedOrTensor] = None,
+    offset: Optional[SeedOrTensor] = None,
     check_nan: bool = False,
 ) -> torch.Tensor:
-    tokens: torch.Tensor = flashinfer.sampling.sampling_from_probs(
-        probs, deterministic=True, generator=generator, check_nan=check_nan
-    )
-    return tokens
-
+    """Categorical sampling from probabilities.
 
-def top_k_top_p_sampling_from_logits_with_generator_op(
-    logits: torch.Tensor,
-    top_k: torch.Tensor,
-    top_p: torch.Tensor,
-    generator: Optional[torch.Generator],
-    check_nan: bool = False,
-) -> torch.Tensor:
-    tokens: torch.Tensor = flashinfer.sampling.top_k_top_p_sampling_from_logits(
-        logits,
-        top_k=top_k,
-        top_p=top_p,
-        filter_apply_order="top_k_first",
+    Randomness: pass ``generator`` (eager) or ``seed``/``offset`` (CUDA graph);
+    see module docstring for the full contract.
+    """
+    tokens: torch.Tensor = flashinfer.sampling.sampling_from_probs(
+        probs,
         deterministic=True,
         check_nan=check_nan,
         generator=generator,
+        seed=seed,
+        offset=offset,
     )
     return tokens
 
 
-def top_k_sampling_from_probs_generator_op(
+def top_k_sampling_from_probs_op(
     probs: torch.Tensor,
     top_k: torch.Tensor,
-    generator: Optional[torch.Generator],
+    *,
+    generator: Optional[torch.Generator] = None,
+    seed: Optional[SeedOrTensor] = None,
+    offset: Optional[SeedOrTensor] = None,
     check_nan: bool = False,
 ) -> torch.Tensor:
+    """Top-k filtered sampling from probabilities.
+
+    Randomness: pass ``generator`` (eager) or ``seed``/``offset`` (CUDA graph);
+    see module docstring for the full contract.
+    """
     tokens: torch.Tensor = flashinfer.sampling.top_k_sampling_from_probs(
         probs,
         top_k=top_k,
         deterministic=True,
         check_nan=check_nan,
         generator=generator,
+        seed=seed,
+        offset=offset,
     )
     return tokens
 
 
-def top_p_sampling_from_probs_generator_op(
+def top_p_sampling_from_probs_op(
     probs: torch.Tensor,
     top_p: torch.Tensor,
-    generator: Optional[torch.Generator],
+    *,
+    generator: Optional[torch.Generator] = None,
+    seed: Optional[SeedOrTensor] = None,
+    offset: Optional[SeedOrTensor] = None,
     check_nan: bool = False,
 ) -> torch.Tensor:
+    """Top-p filtered sampling from probabilities.
+
+    Randomness: pass ``generator`` (eager) or ``seed``/``offset`` (CUDA graph);
+    see module docstring for the full contract.
+    """
     tokens: torch.Tensor = flashinfer.sampling.top_p_sampling_from_probs(
         probs,
         top_p=top_p,
         deterministic=True,
         check_nan=check_nan,
         generator=generator,
+        seed=seed,
+        offset=offset,
     )
     return tokens
 
 
+# The three ops below wrap the mask -> softmax -> renorm pipeline stages 1:1.
+# The wrappers exist so callers stay importable without flashinfer installed
+# (the flashinfer import above is guarded); softmax_op additionally centralizes
+# the PDL env decision.
+
+
+def softmax_op(
+    logits: torch.Tensor,
+    temperature: Optional[torch.Tensor],
+) -> torch.Tensor:
+    probs: torch.Tensor = flashinfer.sampling.softmax(
+        logits, temperature, enable_pdl=get_env_enable_pdl()
+    )
+    return probs
+
+
+def top_k_mask_logits_op(
+    logits: torch.Tensor,
+    top_k: torch.Tensor,
+) -> torch.Tensor:
+    masked: torch.Tensor = flashinfer.sampling.top_k_mask_logits(logits, top_k)
+    return masked
+
+
+def top_p_renorm_probs_op(
+    probs: torch.Tensor,
+    top_p: torch.Tensor,
+) -> torch.Tensor:
+    renormed: torch.Tensor = flashinfer.sampling.top_p_renorm_probs(probs, top_p)
+    return renormed
+
+
 def compute_probs_from_logits_op(
     logits: torch.Tensor,
     temperatures: torch.Tensor,