[TRTLLM-16944][refactor] Simplify TorchSampler strategy grouping and fast path around the fused sampling kernel

pyproject.toml

@@ -873,7 +873,7 @@ show_error_codes = true
 files = [
     "tensorrt_llm/_torch/pyexecutor/sampler/sampler.py",
     "tensorrt_llm/_torch/pyexecutor/sampler/sampler_common.py",
-    "tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py",
+    "tensorrt_llm/_torch/pyexecutor/sampler/sampling_resolution.py",
     "tensorrt_llm/_torch/pyexecutor/sampler/logprobs.py",
     "tensorrt_llm/_torch/pyexecutor/sampler/sampler_features.py",
     "tensorrt_llm/_torch/pyexecutor/sampler/seed_manager.py",
@@ -913,7 +913,7 @@ module = [
 disable_error_code = [
     "no-untyped-call",
     # TODO(TRTLLM-15310): strict mode's warn_redundant_casts flags several cast(...) calls
-    # in sampler_strategy.py (introduced by #16620). Whether a cast is "redundant" is
+    # in the sampler package (introduced by #16620). Whether a cast is "redundant" is
     # binding-dependent -- the full build check (compiled bindings + stubs) flags 2,
     # the bindings-less local/pre-commit check flags ~20 -- so removing casts is a
     # fragile WAR. Suppress redundant-cast for the sampler module until the casts are

tensorrt_llm/_torch/pyexecutor/llm_request.py

@@ -5,8 +5,8 @@
 from copy import copy, deepcopy
 from dataclasses import dataclass, field
 from enum import Enum, auto
-from typing import (TYPE_CHECKING, Any, Container, Dict, Hashable, Iterable,
-                    List, Optional, Union, cast)
+from typing import (Any, Container, Dict, Hashable, Iterable, List, Optional,
+                    Union, cast)
 
 import torch
 
@@ -226,10 +226,6 @@ def pop_all_cache_keys(self) -> List[Hashable]:
         return cache_keys
 
 
-if TYPE_CHECKING:
-    from .sampler.sampler_strategy import Strategy
-
-
 @dataclass(slots=True)
 class PerfTimingInfo:
     """Stores performance timing information for a request."""
@@ -872,8 +868,6 @@ def __init__(
             logprobs_simple_format: bool = False,
             return_routed_experts: bool = False,
             **kwargs):
-        self.py_sampling_strategy: "Strategy | None" = None
-
         self.py_logits_post_processors = kwargs.pop("py_logits_post_processors",
                                                     None)
         # Owned entirely on the Python side: the TorchSampler is the only
@@ -942,8 +936,6 @@ def __init__(
         # Manager-owned helix decode-step counter; see
         # KVCacheManagerV2._set_helix_rank_fields.
         self.py_helix_decode_group_index = 0
-        self.py_draft_logits = None
-        self.py_target_probs = None
         self.py_per_pos_drafted = [0] * MAX_SPEC_DECODE_POSITIONS
         self.py_per_pos_accepted = [0] * MAX_SPEC_DECODE_POSITIONS
         # Cumulative spec-decode counters backing

tensorrt_llm/_torch/pyexecutor/sampler/__init__.py

@@ -15,8 +15,9 @@
 """Sampler package.
 
 The upper-level orchestration (``Sampler`` / ``TorchSampler``)
-lives in ``sampler.py`` and depends on operation-level APIs in
-``sampler_strategy.py``. Implementation-specific kernel providers (FlashInfer,
+lives in ``sampler.py``; ``sampling_resolution.py`` maps a request's sampling
+params onto the fused sampling kernel, for it and for the one-model speculative
+path. Implementation-specific kernel providers (FlashInfer,
 vanilla/PyTorch, TRT-LLM ops) live under ``ops/`` and are selected
 internally, never exposed as interchangeable backends to callers.
 
@@ -29,8 +30,8 @@
                                   │
           ┌───────────┬───────────┼───────────┬───────────┐
           ▼           ▼           ▼           ▼           ▼
-      beam_search  sampler_    logprobs   penalties      ...
-                   strategy                            token_ban
+      beam_search  sampling_   logprobs   penalties      ...
+                   resolution                          token_ban
                                                        top_p_decay
                                                        finish_reasons
                                                        seed_manager
@@ -47,16 +48,15 @@
        ─────────── no intra-package imports ───────────
 
 Feature modules are shown flat, but some depend on others -- notably
-``two_model_spec_dec`` -> ``sampler_strategy`` -> ``beam_search`` ->
-``logprobs`` -> ``sampler_features``. Each feature owns its ``*Store``
+``beam_search`` -> ``logprobs`` -> ``sampler_features``. Each feature owns its ``*Store``
 (persistent per-slot device state) and ``*Handler`` (lifecycle); ``TorchSampler``
 holds one handler per feature and drives them. Types passed *between* modules
 that no single feature owns live in ``sampler_common.py``.
 
 Public symbols from ``sampler.py`` are re-exported here so existing
 ``pyexecutor.sampler`` import paths keep working. The re-export is lazy
 (PEP 562 ``__getattr__``) so that importing lightweight submodules such as
-``pyexecutor.sampler.sampler_strategy`` does not eagerly pull in ``sampler.py``
+``pyexecutor.sampler.sampling_resolution`` does not eagerly pull in ``sampler.py``
 and its heavy dependency chain (which would create import cycles with
 ``speculative.interface``).
 """
@@ -65,7 +65,7 @@
 
 # Submodules of this package — never forward these to sampler.py (that would
 # recurse, since accessing e.g. `.sampler` before it is bound re-enters here).
-_SUBMODULES = frozenset({"sampler", "sampler_common", "sampler_strategy", "ops"})
+_SUBMODULES = frozenset({"sampler", "sampler_common", "sampling_resolution", "ops"})
 
 
 def __getattr__(name: str):

tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py

@@ -43,8 +43,7 @@
 from ..llm_request import LlmRequest, LlmRequestState
 from .logprobs import LogProbsStore
 from .ops.flashinfer import radix_topk_op
-from .ops.vanilla import StrategyMetadata
-from .sampler_common import _get_beam_width_out, int_tensor
+from .sampler_common import _get_beam_width_out, _request_get_sampling_params, int_tensor
 from .sampler_features import _SideStreamCopier
 
 BEAM_SEARCH_PAD_TOKEN = -1
@@ -55,7 +54,7 @@ class BeamSearchEarlyStop(IntEnum):
     ``early_stopping`` (``True`` / ``False`` / ``"never"``).
 
     An ``IntEnum`` so it stays interchangeable with the raw integers coming from
-    ``sampling_config.early_stopping`` and used in the strategy grouping key.
+    ``sampling_config.early_stopping`` and used in ``BeamSearchParams.group_key``.
     Member names follow HF's ``early_stopping`` values (``True`` / ``False`` /
     ``"never"``) directly.
 
@@ -318,7 +317,7 @@ def _gather_beam_path(
 
 
 @dataclass(kw_only=True)
-class BeamSearchMetadata(StrategyMetadata):
+class BeamSearchMetadata:
     """Stateful tensors required by beam_search_sampling_batch_cba."""
 
     cache_indirection: torch.Tensor
@@ -1034,6 +1033,127 @@ def beam_search_sampling_batch_cba(
     return _pad_next_tokens(slot_tok, args.finished_beams.size(1)), softmax
 
 
+class BeamSearchParams(NamedTuple):
+    """A beam-search request's params for one step.
+
+    Not cached per slot like the other sampling params: the beam widths and
+    ``row_stride`` change between the context step and later steps.
+    """
+
+    beam_width_in: int
+    beam_width_out: int
+    # Logits rows the forward laid out per request: the static admission width,
+    # which exceeds beam_width_in while a variable beam width array widens.
+    row_stride: int
+    temperature: float
+    length_penalty: float
+    diversity_rate: float
+    early_stopping: BeamSearchEarlyStop
+
+    @classmethod
+    def from_request(cls, request: LlmRequest) -> "BeamSearchParams":
+        params = _request_get_sampling_params(request)
+        assert params.beam_width_in is not None and params.beam_width_out is not None, (
+            "beam_width_in and beam_width_out must be specified for beam search"
+        )
+        temperature = params.temperature or 1.0
+        assert temperature > 0, f"beam search needs temperature > 0, got {temperature}"
+        return cls(
+            beam_width_in=params.beam_width_in,
+            beam_width_out=params.beam_width_out,
+            row_stride=params.row_stride or params.beam_width_in,
+            temperature=temperature,
+            length_penalty=params.length_penalty or 0.0,
+            diversity_rate=params.beam_search_diversity_rate or 0.0,
+            early_stopping=BeamSearchEarlyStop.from_raw(params.early_stopping),
+        )
+
+    @property
+    def group_key(self) -> tuple[int, int, int, BeamSearchEarlyStop]:
+        """What requests sampled in one beam-search step must share.
+
+        ``row_stride`` is part of it: a context request (one logits row) and a
+        generation request (``py_beam_width`` rows) can have the same beam
+        widths under different beam width arrays.
+        """
+        return self.beam_width_in, self.beam_width_out, self.row_stride, self.early_stopping
+
+
+class BeamSearchStep:
+    """One candidate-beams-array step for a group of beam-search requests.
+
+    Every early_stopping mode runs here; the mode only selects the done verdict
+    computed inside :func:`beam_search_sampling_batch_cba`.
+    """
+
+    def __init__(
+        self,
+        beam_width_in: int,
+        beam_width_out: int,
+        row_stride: int,
+        temperature: torch.Tensor,
+        length_penalty: Optional[torch.Tensor],
+        diversity_rate: Optional[torch.Tensor],
+        early_stopping: BeamSearchEarlyStop,
+        *,
+        computes_probs: bool = False,
+    ):
+        self._beam_width_in = beam_width_in
+        self._beam_width_out = beam_width_out
+        self._row_stride = row_stride
+        self.temperature = temperature
+        self._length_penalty = length_penalty
+        self._diversity_rate = diversity_rate
+        self._early_stopping = early_stopping
+        self._computes_probs = computes_probs
+
+    @classmethod
+    def from_params(
+        cls, params: list[BeamSearchParams], device: torch.device, *, return_probs: bool
+    ) -> "BeamSearchStep":
+        (group_key,) = set(p.group_key for p in params)
+        beam_width_in, beam_width_out, row_stride, early_stopping = group_key
+
+        def _to_device(values: list[float]) -> torch.Tensor:
+            return torch.tensor(values, dtype=torch.float32, pin_memory=prefer_pinned()).to(
+                device=device, non_blocking=True
+            )
+
+        # Unset length_penalty / diversity_rate stay None so the op skips them.
+        length_penalties = [p.length_penalty for p in params]
+        diversity_rates = [p.diversity_rate for p in params]
+        return cls(
+            beam_width_in,
+            beam_width_out,
+            row_stride,
+            _to_device([p.temperature for p in params]),
+            _to_device(length_penalties) if any(length_penalties) else None,
+            _to_device(diversity_rates) if any(diversity_rates) else None,
+            early_stopping,
+            computes_probs=return_probs,
+        )
+
+    def sample(
+        self, logits: torch.Tensor, *, metadata: BeamSearchMetadata
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+        """Sample the group's logits rows, ``row_stride`` per request."""
+        # Temperature is applied before the op slices the padding rows off, so
+        # it must cover every row the forward laid out.
+        temperature = self.temperature.repeat_interleave(self._row_stride).unsqueeze(-1)
+        return beam_search_sampling_batch_cba(
+            logits / temperature,
+            beam_width_in=self._beam_width_in,
+            beam_width_out=self._beam_width_out,
+            row_stride=self._row_stride,
+            beam_search_args=metadata,
+            temperature=None,
+            early_stopping=self._early_stopping,
+            length_penalty=self._length_penalty,
+            diversity_rate=self._diversity_rate,
+            return_probs=self._computes_probs,
+        )
+
+
 BeamHistoryBuilder: TypeAlias = Callable[[], BeamHistory | None]
 """Builder for BeamHistory.
 
@@ -1374,8 +1494,8 @@ def build_metadata(
         store = self._store
         assert store is not None
         assert seq_lens is not None, "seq_lens is required for beam search"
-        # Reuse the precomputed CUDA tensors when the strategy group
-        # covers the full batch (typical single-strategy case);
+        # Reuse the precomputed CUDA tensors when the group covers the full
+        # batch (the typical case);
         # otherwise fall back to a per-group H2D for the subset.
         if group_req_indices.size(0) == num_requests:
             group_seq_slots_cuda = seq_slots_cuda

tensorrt_llm/_torch/pyexecutor/sampler/logprobs.py

@@ -446,7 +446,7 @@ def _gather_src_dst_indices(
                 ).squeeze(-1)  # flattened (step, slot)
 
                 # 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 FlashInferGroupedStrategySampler, s.t., e.g.,
+                # NB: Computation of sampled rank could be lowered into the per-group sampling, 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(