[pull] main from NVIDIA:main

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/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

@@ -38,6 +38,7 @@
 import numpy as np
 import torch
 
+from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE
 from tensorrt_llm._torch.pyexecutor.make_decoding_batch_input_output import (
     MakeDecodingBatchInputOutput,
 )
@@ -87,11 +88,11 @@
 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(