[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
(TorchSampler.__init__, SpecWorkerBase) with a clear ImportError, so
sampling_utils stays importable for torch-only consumers (AutoDeploy demollm,
drafting greedy helpers).

- IS_FLASHINFER_AVAILABLE semantics are unchanged (its ~46 legitimate fallback
  consumers are not touched).
- ops/flashinfer.py: keep flashinfer.sampling behind an import guard so the
  module stays importable without flashinfer; ops are only reached after the
  hard-dependency check at construction time.
- TorchSampler.__init__ enforces the hard dependency once (out of the
  CUDA-graph-captured loop) before selecting the grouped sampler class.
- Flatten the single-implementation grouped sampler: drop the
  GroupedStrategySampler ABC and keep only FlashInferGroupedStrategySampler.
- 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

@@ -70,7 +70,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

@@ -2436,7 +2436,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,
 ):
@@ -2452,7 +2451,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,
@@ -2471,7 +2469,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)
@@ -2483,7 +2480,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

@@ -796,7 +796,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

@@ -14,9 +14,8 @@
 
 """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).
+These ops depend on flashinfer; the import is guarded so the module stays
+importable without it.
 """
 
 from typing import Optional

tensorrt_llm/_torch/pyexecutor/sampler/ops/interface.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/ops/vanilla.py

@@ -84,14 +84,15 @@
 from ...speculative.spec_tree_manager import SpecTreeManager
 from ...utils import torch_multi_arange
 from ..finish_reason import FinishedState
+from ..flashinfer_utils import IS_FLASHINFER_AVAILABLE
 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 .sampling_utils import (
     BEAM_SEARCH_PAD_TOKEN,
     GREEDY,
     BeamSearchMetadata,
+    FlashInferGroupedStrategySampler,
     GenericStrategyKeyType,
     Strategy,
     StrategyMetadata,
@@ -2321,7 +2322,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
 
@@ -2345,13 +2345,15 @@ 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,
-            ),
-        )
+        # The Torch sampler hard-depends on flashinfer. Enforce it once here, at
+        # construction, so the check stays out of the CUDA-graph-captured
+        # sampling loop.
+        if not IS_FLASHINFER_AVAILABLE:
+            raise ImportError(
+                "flashinfer is not available, please install the version pinned "
+                "in requirements.txt."
+            )
+        self._grouped_sampler_cls = FlashInferGroupedStrategySampler
 
         # AutoDeploy build creates the sampler in inference mode,
         # which would disallow in-place mutating of new_tokens.
@@ -4598,7 +4600,7 @@ def _process_logprobs(
             sampled_indices_cuda = group_next_tokens_cuda.squeeze(1)
 
             # sampled_rank_cuda contains the 0-based rank, it will be corrected to 1-based in handle_logprobs
-            # NB: Computation of sampled rank could be lowered into GroupedStrategySampler, s.t., e.g., for
+            # NB: Computation of sampled rank could be lowered into FlashInferGroupedStrategySampler, s.t., e.g., for
             #     greedy sampling, logits management and log_softmax could be completely skipped (sampled rank
             #     computation is trivial in this case).
             sampled_rank_cuda = _Fusions.determine_sampled_rank(