[TRTLLM-13212][refactor] Make FlashInfer a hard dependency for Torch sampler and spec

Convert the Torch sampler and one-model speculative decoding from
"FlashInfer-optional with torch fallback" to a hard FlashInfer dependency.
The failure is raised at sampler / speculative-worker construction time
(resolve_sampling_backend, SpecWorkerBase) via a new require_flashinfer
helper and a sentinel-based bind_flashinfer_sampling, so sampling_utils
stays importable for torch-only consumers (AutoDeploy demollm, drafting
greedy helpers).

- Add require_flashinfer() in flashinfer_utils.py; IS_FLASHINFER_AVAILABLE
  semantics are unchanged (its ~46 legitimate fallback consumers are not
  touched).
- ops/flashinfer.py: drop the import guard, bind flashinfer.sampling at
  construction via a sentinel; compiled bodies do no import work.
- Remove the SimpleGroupedStrategySampler torch backend and the
  disable_flashinfer_sampling LlmArg + its plumbing; update the api_stability
  and usage golden files and the sampling/telemetry docs.
- compute_probs_from_logits becomes FlashInfer-only (drop the CUDA C++ op and
  CPU vanilla branches).
- speculative/interface.py: drop the availability guards, keep the
  flashinfer>=0.6.4 version gate and its version-gated branches.
- Delete now-dead fallback kernels (forward_native_sampling,
  _apply_top_k_top_p, _random_sample, vanilla.compute_probs_from_logits_op);
  keep beam-search/greedy/rejection/_Fusions and the torch sample() path used
  by rejection-sampling draft probs and demollm.
- Prune torch-vs-flashinfer parametrized sampler tests.

Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>

Author: zhaoyangwang-nvidia

docs/source/developer-guide/telemetry.md

@@ -63,7 +63,6 @@ unset or when the safety sanitizer rejects the runtime value.
 | `cuda_graph_config.mode` | `Literal['decode']` | `categorical` |  | `decode`, `encode` |
 | `cuda_graph_config.num_tokens` | `Optional[List[Annotated[int, Gt(gt=0)]]]` | `value` |  |  |
 | `cuda_graph_config.seq_lens` | `Optional[List[Annotated[int, Gt(gt=0)]]]` | `value` |  |  |
-| `disable_flashinfer_sampling` | `<class 'bool'>` | `value` |  |  |
 | `disable_overlap_scheduler` | `<class 'bool'>` | `value` |  |  |
 | `dtype` | `<class 'str'>` | `categorical` | allowlist | `auto`, `float16`, `bfloat16`, `float32` |
 | `dwdp_config.contention_opt` | `<class 'bool'>` | `value` |  |  |

docs/source/features/sampling.md

@@ -104,14 +104,12 @@ llm.generate(["Hello, my name is",
 ### Performance
 
 The Torch Sampler leverages the optimized sampling kernels provided by
-[FlashInfer](https://docs.flashinfer.ai/api/sampling.html). The sampler
-also uses the [sorting-free implementations](https://flashinfer.ai/2025/03/10/sampling.html)
+[FlashInfer](https://docs.flashinfer.ai/api/sampling.html), which is a required
+dependency for the Torch Sampler. The sampler also uses the
+[sorting-free implementations](https://flashinfer.ai/2025/03/10/sampling.html)
 whenever possible. This optimization does not compute the complete set of token sampling probabilities
 (after top-k / top-p masking etc.), which typically can be omitted unless requested by the user or
 required for speculative decoding (rejection sampling).
-In case of unexpected problems, the use of FlashInfer in Torch Sampler can
-be disabled via the `disable_flashinfer_sampling` config option (note that this option is likely
-to be removed in a future TensorRT LLM release).
 
 Moreover, Torch Sampler internally batches requests with compatible sampling parameters. This
 can greatly reduce the overall latency of the sampling step when request batches are comprised

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -2343,7 +2343,6 @@ def create_torch_sampler_args(
     speculative_config: SpeculativeConfig,
     max_beam_width: int,
     disable_overlap_scheduler: bool,
-    disable_flashinfer_sampling: bool,
     enable_async_worker: bool,
     enable_speculative_beam_history_d2h: bool,
 ):
@@ -2359,7 +2358,6 @@ def create_torch_sampler_args(
         max_total_draft_tokens=max_total_draft_tokens,
         max_num_sequences=max_num_sequences,
         max_beam_width=max_beam_width,
-        disable_flashinfer_sampling=disable_flashinfer_sampling,
         disable_overlap_scheduler=disable_overlap_scheduler,
         enable_async_worker=enable_async_worker,
         enable_speculative_beam_history_d2h=enable_speculative_beam_history_d2h,
@@ -2378,7 +2376,6 @@ def instantiate_sampler(
     speculative_config: SpeculativeConfig,
     decoding_config: trtllm.DecodingConfig,
     kv_cache_config: KvCacheConfig,
-    disable_flashinfer_sampling: bool,
 ):
     enable_async_worker = (confidential_compute_enabled()
                            or llm_args.sampler_force_async_worker)
@@ -2390,7 +2387,6 @@ def instantiate_sampler(
         speculative_config=speculative_config,
         max_beam_width=max_beam_width,
         disable_overlap_scheduler=llm_args.disable_overlap_scheduler,
-        disable_flashinfer_sampling=disable_flashinfer_sampling,
         enable_async_worker=enable_async_worker,
         enable_speculative_beam_history_d2h=llm_args.
         enable_speculative_beam_history_d2h,

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -763,7 +763,6 @@ def drafting_loop_wrapper(model):
             speculative_config=spec_config,
             decoding_config=decoding_config,
             kv_cache_config=kv_cache_config,
-            disable_flashinfer_sampling=llm_args.disable_flashinfer_sampling,
         )
         logger.info(f"Using Sampler: {type(sampler).__name__}")
 

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

@@ -15,11 +15,19 @@
 """FlashInfer-accelerated sampling kernels.
 
 Pure kernel functions with no dependency on the sampling_utils interface
-or other backend implementation modules. All flashinfer imports are guarded by
-IS_FLASHINFER_AVAILABLE. Beam search is excluded (torch-only per design).
+or other backend implementation modules. Beam search is excluded (torch-only
+per design).
+
+flashinfer is a hard dependency for these ops. The module stays importable
+without flashinfer (so ``sampling_utils`` and its torch-only ``greedy`` helper
+remain importable for AutoDeploy / drafting consumers), but the actual
+``flashinfer.sampling`` module is bound at sampler / speculative-worker
+construction time via ``bind_flashinfer_sampling`` -- which raises a clear
+``ImportError`` when flashinfer is absent. Calling these ops before binding
+raises a descriptive ``RuntimeError``.
 """
 
-from typing import Optional
+from typing import Any, NoReturn, Optional
 
 import torch
 
@@ -29,14 +37,45 @@
     import flashinfer.sampling
 
 
+class _UnboundFlashInferSampling:
+    """Sentinel used until ``bind_flashinfer_sampling`` binds the real module."""
+
+    def __getattr__(self, name: str) -> NoReturn:
+        raise RuntimeError(
+            "FlashInfer sampling ops were called before flashinfer was bound. "
+            "Construct the Torch sampler (resolve_sampling_backend) or the "
+            "speculative worker before invoking these ops."
+        )
+
+
+# Bound to ``flashinfer.sampling`` by ``bind_flashinfer_sampling`` at construction.
+_flashinfer_sampling: Any = _UnboundFlashInferSampling()
+
+
+def bind_flashinfer_sampling(feature_name: str = "Torch sampler") -> None:
+    """Bind the real ``flashinfer.sampling`` module, requiring flashinfer.
+
+    Idempotent-friendly: re-binding to the same module is harmless. Raises
+    ``ImportError`` when flashinfer is unavailable.
+    """
+    if not IS_FLASHINFER_AVAILABLE:
+        raise ImportError(
+            f"{feature_name} requires flashinfer, but flashinfer is not "
+            "available on this platform. flashinfer is a hard dependency here; "
+            "please install the version pinned in requirements.txt."
+        )
+    global _flashinfer_sampling
+    _flashinfer_sampling = flashinfer.sampling
+
+
 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,
 ) -> torch.Tensor:
-    tokens: torch.Tensor = flashinfer.sampling.top_k_top_p_sampling_from_logits(
+    tokens: torch.Tensor = _flashinfer_sampling.top_k_top_p_sampling_from_logits(
         logits, top_k, top_p, seed=seed, offset=offset
     )
     return tokens
@@ -47,7 +86,7 @@ def sampling_from_probs_op(
     seed: Optional[torch.Tensor] = None,
     offset: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
-    tokens: torch.Tensor = flashinfer.sampling.sampling_from_probs(
+    tokens: torch.Tensor = _flashinfer_sampling.sampling_from_probs(
         probs, deterministic=True, seed=seed, offset=offset
     )
     return tokens
@@ -57,7 +96,7 @@ def softmax_op(
     logits: torch.Tensor,
     temperature: Optional[torch.Tensor],
 ) -> torch.Tensor:
-    probs: torch.Tensor = flashinfer.sampling.softmax(
+    probs: torch.Tensor = _flashinfer_sampling.softmax(
         logits, temperature, enable_pdl=get_env_enable_pdl()
     )
     return probs
@@ -67,15 +106,15 @@ 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)
+    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)
+    renormed: torch.Tensor = _flashinfer_sampling.top_p_renorm_probs(probs, top_p)
     return renormed
 
 
@@ -84,7 +123,7 @@ def sampling_from_probs_generator_op(
     generator: Optional[torch.Generator],
     check_nan: bool = False,
 ) -> torch.Tensor:
-    tokens: torch.Tensor = flashinfer.sampling.sampling_from_probs(
+    tokens: torch.Tensor = _flashinfer_sampling.sampling_from_probs(
         probs, deterministic=True, generator=generator, check_nan=check_nan
     )
     return tokens
@@ -97,7 +136,7 @@ def top_k_top_p_sampling_from_logits_with_generator_op(
     generator: Optional[torch.Generator],
     check_nan: bool = False,
 ) -> torch.Tensor:
-    tokens: torch.Tensor = flashinfer.sampling.top_k_top_p_sampling_from_logits(
+    tokens: torch.Tensor = _flashinfer_sampling.top_k_top_p_sampling_from_logits(
         logits,
         top_k=top_k,
         top_p=top_p,
@@ -115,7 +154,7 @@ def top_k_sampling_from_probs_generator_op(
     generator: Optional[torch.Generator],
     check_nan: bool = False,
 ) -> torch.Tensor:
-    tokens: torch.Tensor = flashinfer.sampling.top_k_sampling_from_probs(
+    tokens: torch.Tensor = _flashinfer_sampling.top_k_sampling_from_probs(
         probs,
         top_k=top_k,
         deterministic=True,
@@ -131,7 +170,7 @@ def top_p_sampling_from_probs_generator_op(
     generator: Optional[torch.Generator],
     check_nan: bool = False,
 ) -> torch.Tensor:
-    tokens: torch.Tensor = flashinfer.sampling.top_p_sampling_from_probs(
+    tokens: torch.Tensor = _flashinfer_sampling.top_p_sampling_from_probs(
         probs,
         top_p=top_p,
         deterministic=True,
@@ -154,10 +193,10 @@ def compute_probs_from_logits_op(
     optional top-k masking (via fused flashinfer softmax+temp).
     """
     if top_k is not None:
-        logits = flashinfer.sampling.top_k_mask_logits(logits, top_k)
-    probs: torch.Tensor = flashinfer.sampling.softmax(
+        logits = _flashinfer_sampling.top_k_mask_logits(logits, top_k)
+    probs: torch.Tensor = _flashinfer_sampling.softmax(
         logits, temperatures, enable_pdl=get_env_enable_pdl()
     )
     if top_p is not None:
-        probs = flashinfer.sampling.top_p_renorm_probs(probs, top_p)
+        probs = _flashinfer_sampling.top_p_renorm_probs(probs, top_p)
     return probs

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

@@ -14,44 +14,30 @@
 
 """Sampling strategy-grouping backend selection.
 
-Holds ``SamplerConfig`` and the ``resolve_sampling_backend`` factory that picks
-the strategy-grouping sampler (FlashInfer vs simple/torch) once at init time, so
-no per-call dispatch happens inside CUDA graph capture. The three grouping
-callables are tightly coupled (they share the same key type), so the whole
-sampler class is bound as a unit rather than field by field.
+Holds the ``resolve_sampling_backend`` factory that binds the strategy-grouping
+sampler once at init time, so no per-call dispatch happens inside CUDA graph
+capture. flashinfer is a hard dependency for the Torch sampler: the factory
+raises at sampler construction time if flashinfer is unavailable.
 """
 
 from collections.abc import Hashable
-from dataclasses import dataclass
 from typing import Type, cast
 
-from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE
+from tensorrt_llm._torch.pyexecutor.sampler.ops.flashinfer import bind_flashinfer_sampling
 from tensorrt_llm._torch.pyexecutor.sampler.sampling_utils import (
     FlashInferGroupedStrategySampler,
     GroupedStrategySampler,
-    SimpleGroupedStrategySampler,
 )
 
 
-@dataclass(frozen=True, kw_only=True)
-class SamplerConfig:
-    """Configuration used to resolve the sampling backend at init time."""
+def resolve_sampling_backend() -> Type[GroupedStrategySampler[Hashable]]:
+    """Bind the strategy-grouping sampler class at init time; CUDA-graph safe.
 
-    use_flashinfer: bool = False
-
-
-def resolve_sampling_backend(
-    is_cuda: bool,
-    config: SamplerConfig,
-) -> Type[GroupedStrategySampler[Hashable]]:
-    """Pick the strategy-grouping sampler class at init time; CUDA-graph safe.
-
-    Selection order:
-      1. FlashInfer  — is_cuda AND IS_FLASHINFER_AVAILABLE AND config.use_flashinfer
-      2. Torch       — everything else (including CPU)
+    The Torch sampler hard-depends on flashinfer, so this binds the flashinfer
+    sampling module (raising a clear error at construction time when flashinfer
+    is unavailable) rather than silently falling back to a torch implementation.
     """
-    # The two samplers use different key types (Strategy vs a narrower FlashInfer
-    # key); the caller only forwards the grouping callables, so erase to Hashable.
-    if is_cuda and IS_FLASHINFER_AVAILABLE and config.use_flashinfer:
-        return cast(Type[GroupedStrategySampler[Hashable]], FlashInferGroupedStrategySampler)
-    return cast(Type[GroupedStrategySampler[Hashable]], SimpleGroupedStrategySampler)
+    bind_flashinfer_sampling("Torch sampler")
+    # FlashInferGroupedStrategySampler uses a narrower FlashInfer key type; the
+    # caller only forwards the grouping callables, so erase to Hashable.
+    return cast(Type[GroupedStrategySampler[Hashable]], FlashInferGroupedStrategySampler)

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

@@ -353,71 +353,6 @@ def _safely_apply_temperature_inplace(
     return logits_inout.div_(safe_temp.unsqueeze(dim=1))
 
 
-def _apply_top_k_top_p(
-    logits: torch.Tensor,
-    k: Optional[torch.Tensor],
-    p: Optional[torch.Tensor],
-) -> torch.Tensor:
-    logits_sort, logits_idx = logits.sort(dim=-1, descending=False)
-    if k is not None:
-        top_k_mask = logits_sort.size(1) - k.to(torch.long)
-        top_k_mask = top_k_mask.clamp(min=0)
-        top_k_mask = logits_sort.gather(1, top_k_mask.unsqueeze(dim=1))
-        top_k_mask = logits_sort < top_k_mask
-        logits_sort.masked_fill_(top_k_mask, -float("inf"))
-    if p is not None:
-        probs_sort = logits_sort.softmax(dim=-1)
-        probs_sum = torch.cumsum(probs_sort, dim=-1, out=probs_sort)
-        top_p_mask = probs_sum <= 1 - p.unsqueeze(dim=1)
-        top_p_mask[:, -1] = False
-        logits_sort.masked_fill_(top_p_mask, -float("inf"))
-    return logits_sort.scatter(dim=-1, index=logits_idx, src=logits_sort)
-
-
-def _random_sample(probs: torch.Tensor) -> torch.Tensor:
-    q = torch.empty_like(probs).exponential_()
-    return probs.div_(q).argmax(dim=-1).view(-1)
-
-
-def forward_native_sampling(
-    logits: torch.Tensor,
-    k: Optional[torch.Tensor],
-    p: Optional[torch.Tensor],
-) -> torch.Tensor:
-    logits = _apply_top_k_top_p(logits, k, p)
-    probs = logits.softmax(dim=-1, dtype=torch.float32)
-    return _random_sample(probs)
-
-
-def compute_probs_from_logits_op(
-    logits: torch.Tensor,
-    temperatures: torch.Tensor,
-    top_k: Optional[torch.Tensor],
-    top_p: Optional[torch.Tensor],
-) -> torch.Tensor:
-    """Pure-PyTorch CPU fallback for probability computation."""
-    is_greedy = temperatures <= _GREEDY_TEMPERATURE_THRESHOLD
-    # Greedy rows must pick the argmax of the *original* logits (before temperature).
-    # Capture the argmax up front; _safely_apply_temperature_inplace then guards the
-    # division against the greedy sentinel.
-    argmax_ids = logits.argmax(dim=-1, keepdim=True)
-
-    logits = _safely_apply_temperature_inplace(logits, temperatures)
-    logits = _apply_top_k_top_p(logits, top_k, top_p)
-    probs = logits.softmax(dim=-1, dtype=torch.float32)
-
-    # Turn the greedy rows into a one-hot at argmax by editing `probs` in place,
-    # instead of building a full [batch, vocab] one-hot buffer and a [batch, vocab]
-    # torch.where copy. The torch.where here only runs on a [batch, 1] tensor.
-    # NB: argwhere/index-select on the greedy rows would give a data-dependent shape
-    # that breaks the surrounding torch.compile graph, so we keep it dense.
-    greedy_col = is_greedy.unsqueeze(1)
-    new_at_argmax = torch.where(greedy_col, 1.0, probs.gather(1, argmax_ids))
-    probs.masked_fill_(greedy_col, 0.0)
-    probs.scatter_(1, argmax_ids, new_at_argmax)
-    return probs
-
-
 class _Fusions:
     @staticmethod
     @torch.compile(dynamic=None, fullgraph=True)

tensorrt_llm/_torch/pyexecutor/sampler/sampler.py

@@ -87,7 +87,7 @@
 from ..llm_request import LlmRequest, LlmRequestState, get_draft_token_length
 from ..resource_manager import ResourceManager, ResourceManagerType
 from ..scheduler import ScheduledRequests
-from .ops.interface import SamplerConfig, resolve_sampling_backend
+from .ops.interface import resolve_sampling_backend
 from .sampling_utils import (
     BEAM_SEARCH_PAD_TOKEN,
     GREEDY,
@@ -2323,7 +2323,6 @@ class Args:
         max_beam_width: int
         max_total_draft_tokens: int
         disable_overlap_scheduler: bool = False
-        disable_flashinfer_sampling: bool = False
         enable_async_worker: bool = False
         enable_speculative_beam_history_d2h: bool = False
 
@@ -2347,13 +2346,7 @@ def __init__(self, args: Args):
         self.LOGPROBS_SHAPE = (self.max_num_sequences, self.max_beam_width, self.max_tokens)
         self.TOPK_LOGPROBS_SHAPE = (self.max_num_sequences, self.max_tokens, self.max_topk_logprobs)
 
-        self._grouped_sampler_cls = resolve_sampling_backend(
-            is_cuda=True,
-            config=SamplerConfig(
-                # IS_FLASHINFER_AVAILABLE is checked inside resolve_sampling_backend.
-                use_flashinfer=not args.disable_flashinfer_sampling,
-            ),
-        )
+        self._grouped_sampler_cls = resolve_sampling_backend()
 
         # AutoDeploy build creates the sampler in inference mode,
         # which would disallow in-place mutating of new_tokens.