[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,8 +15,13 @@
 """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, but the import is guarded so the
+module stays importable without flashinfer (keeping ``sampling_utils`` importable
+for its torch-only consumers). These ops are only reached after the Torch sampler
+or a speculative worker has verified flashinfer at construction time.
 """
 
 from typing import Optional

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

@@ -14,44 +14,34 @@
 
 """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 picks the strategy-grouping
+sampler once at init time, so no per-call dispatch happens inside CUDA graph
+capture.
 """
 
 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.sampling_utils import (
     FlashInferGroupedStrategySampler,
     GroupedStrategySampler,
-    SimpleGroupedStrategySampler,
 )
 
 
-@dataclass(frozen=True, kw_only=True)
-class SamplerConfig:
-    """Configuration used to resolve the sampling backend at init time."""
-
-    use_flashinfer: bool = False
-
-
-def resolve_sampling_backend(
-    is_cuda: bool,
-    config: SamplerConfig,
-) -> Type[GroupedStrategySampler[Hashable]]:
+def resolve_sampling_backend() -> 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 raises 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)
+    if not IS_FLASHINFER_AVAILABLE:
+        raise ImportError(
+            "Torch sampler requires flashinfer, but flashinfer is not available "
+            "on this platform. flashinfer is a hard dependency here; please "
+            "install the version pinned in requirements.txt."
+        )
+    # 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.

tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py

@@ -26,14 +26,10 @@
 
 import torch
 
-from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE
 from tensorrt_llm._torch.pyexecutor.sampler.ops import flashinfer, vanilla
 
-# NB: these flashinfer op wrappers are plain Python functions that are safe to
-# import even without flashinfer installed (the flashinfer import inside
-# ops/flashinfer.py is itself guarded); they are only *called* under
-# IS_FLASHINFER_AVAILABLE. Importing them unconditionally keeps them defined for
-# static analysis (they are referenced unconditionally in the strategy impls).
+# These op wrappers are safe to import without flashinfer installed; they are
+# only called on the flashinfer sampler / speculative-worker paths.
 from tensorrt_llm._torch.pyexecutor.sampler.ops.flashinfer import (
     sampling_from_probs_generator_op as sampling_from_probs_generator_op,
 )
@@ -179,6 +175,8 @@ def sample(
     group_metadata: StrategyMetadata | None = None,
     return_probs: bool = True,
 ) -> tuple[torch.Tensor, torch.Tensor | None, float | None]:
+    """Torch-native single-strategy sampling for the rejection-sampling
+    draft-probs path (sampler.py), independent of the flashinfer sampler."""
     softmax: torch.Tensor | None
     # 'cast' needed b/c of https://github.com/python/mypy/issues/19081
     match strategy:
@@ -266,57 +264,6 @@ def sample_grouped_strategies(
         raise NotImplementedError
 
 
-class SimpleGroupedStrategySampler(GroupedStrategySampler[Strategy]):
-    STRATEGY_KEY_TYPE: TypeAlias = Strategy
-
-    @override
-    @staticmethod
-    def strategy_grouping_key(strategy: Strategy) -> STRATEGY_KEY_TYPE:
-        return strategy
-
-    @override
-    @staticmethod
-    def get_metadata_type_for_group(
-        strategy_key: STRATEGY_KEY_TYPE,
-    ) -> Type[StrategyMetadata] | None:
-        match strategy_key:
-            case ("beam_search", _, _, _):
-                return BeamSearchMetadata
-            case _:
-                return None
-
-    @override
-    @staticmethod
-    def sample_grouped_strategies(
-        group_key: STRATEGY_KEY_TYPE,
-        strategies: list[Strategy],
-        logits: torch.Tensor,
-        *,
-        group_logit_indices: torch.Tensor | None = None,
-        generator: torch.Generator | None = None,
-        return_probs: bool,
-        group_metadata: StrategyMetadata | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor | None, float | torch.Tensor | None]:
-        if group_key[0] == "beam_search":
-            beam_width_in = group_key[1]
-        else:
-            beam_width_in = 1
-
-        if group_logit_indices is not None:
-            logits = logits[group_logit_indices]
-        assert logits.size(0) == beam_width_in * len(strategies)
-
-        assert all(strategy == group_key for strategy in strategies), "group must be consistent"
-
-        return sample(
-            group_key,
-            logits,
-            generator=generator,
-            return_probs=return_probs,
-            group_metadata=group_metadata,
-        )
-
-
 class _StrategyImpls:
     class StrategyImpl(abc.ABC):
         @classmethod
@@ -928,24 +875,15 @@ def compute_probs_from_logits(
     top_k: Optional[torch.Tensor],
     top_p: Optional[torch.Tensor],
 ) -> torch.Tensor:
-    """Compute filtered+normalized probs. Dispatches: flashinfer → C++ op → CPU.
+    """Compute filtered+normalized probs via flashinfer (hard dependency).
 
-    ``temperatures``, ``top_k``, ``top_p`` are per-request tensors matching
-    the spec-decoding call site in interface.py.
+    ``temperatures``, ``top_k``, ``top_p`` are per-request tensors matching the
+    spec-decoding call site in interface.py.
     """
     if top_k is not None:
         top_k = sanitize_top_k(top_k, logits.shape[-1])
 
-    if logits.is_cuda and IS_FLASHINFER_AVAILABLE:
-        return flashinfer.compute_probs_from_logits_op(logits, temperatures, top_k, top_p)
-    if logits.is_cuda:
-        # TRT-LLM C++ op (CUDA, no flashinfer). The op keeps a skip_temperature
-        # flag; temperature is always applied here.
-        probs: torch.Tensor = torch.ops.trtllm.compute_probs_from_logits_op(
-            logits, temperatures, top_k, top_p, False
-        )
-        return probs
-    return vanilla.compute_probs_from_logits_op(logits, temperatures, top_k, top_p)
+    return flashinfer.compute_probs_from_logits_op(logits, temperatures, top_k, top_p)
 
 
 @torch.compile(options={"max-autotune": True})
@@ -968,12 +906,9 @@ def sampling_batch_spec_dec_one_model(
     is_greedy = temperatures <= vanilla._GREEDY_TEMPERATURE_THRESHOLD
     greedy_tokens = logits.argmax(dim=-1)
     logits = vanilla._safely_apply_temperature_inplace(logits, temperatures)
-    if IS_FLASHINFER_AVAILABLE:
-        sampled = flashinfer.top_k_top_p_sampling_from_logits_op(
-            logits, top_k, top_p, seed=seed, offset=offset
-        )
-    else:
-        sampled = vanilla.forward_native_sampling(logits, top_k, top_p)
+    sampled = flashinfer.top_k_top_p_sampling_from_logits_op(
+        logits, top_k, top_p, seed=seed, offset=offset
+    )
     # argmax yields int64; cast so torch.where preserves the sampler's dtype
     # (flashinfer returns int32) instead of promoting the result to int64.
     return torch.where(is_greedy, greedy_tokens.to(sampled.dtype), sampled)
@@ -989,14 +924,8 @@ def sampling_batch_spec_dec_one_model_for_rejection(
     offset: Optional[torch.Tensor] = None,
 ) -> tuple[torch.Tensor, torch.Tensor]:
     """Draft sampler returning tokens AND probs for the downstream rejection-sampling path."""
+    # Rejection sampling relies on flashinfer's seed/offset support for
+    # determinism and cross-rank consistency.
     probs = compute_probs_from_logits(logits, temperatures, top_k, top_p)
-    if not IS_FLASHINFER_AVAILABLE:
-        # The torch-native fallback samples from the global RNG and ignores
-        # seed/offset, which breaks determinism and cross-rank consistency that
-        # one-model speculative rejection sampling relies on. Require flashinfer
-        # instead of silently degrading (matches the pre-refactor behavior).
-        raise RuntimeError(
-            "Rejection sampling for one-model speculative decoding requires flashinfer"
-        )
     tokens = flashinfer.sampling_from_probs_op(probs, seed=seed, offset=offset)
     return tokens, probs