diff --git a/docs/source/developer-guide/telemetry.md b/docs/source/developer-guide/telemetry.md index e2724a9ed790..e5b763c19386 100644 --- a/docs/source/developer-guide/telemetry.md +++ b/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` | `` | `value` | | | | `disable_overlap_scheduler` | `` | `value` | | | | `dtype` | `` | `categorical` | allowlist | `auto`, `float16`, `bfloat16`, `float32` | | `dwdp_config.contention_opt` | `` | `value` | | | diff --git a/docs/source/features/sampling.md b/docs/source/features/sampling.md index fc3cb3ed9324..d5a9eaef0685 100644 --- a/docs/source/features/sampling.md +++ b/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 diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 48d0c0dccdc4..cda3d9a44136 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/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, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index ff0c626d70f7..2ef732c166b4 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/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__}") diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py index 8a503b18f338..902472d8f562 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py +++ b/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 diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/interface.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/interface.py deleted file mode 100644 index 9fb70e493bee..000000000000 --- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/interface.py +++ /dev/null @@ -1,57 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""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. -""" - -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]]: - """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 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) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py index b01406cb8668..ffb962a60fc6 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py +++ b/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) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 1ad625d6d234..e7fe112ace6f 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/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( diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py index fb661cb6f496..0c4d59b0ecc4 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py @@ -22,18 +22,14 @@ import sys from collections.abc import Hashable from dataclasses import dataclass -from typing import Any, Generic, Literal, Optional, Type, TypeAlias, TypeVar, cast +from typing import Any, Literal, Optional, Type, TypeAlias, TypeVar, cast 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, ) @@ -231,92 +227,6 @@ def sample( GenericStrategyKeyType = TypeVar("GenericStrategyKeyType", bound=Hashable) -class GroupedStrategySampler(Generic[GenericStrategyKeyType], abc.ABC): - @staticmethod - @abc.abstractmethod - def strategy_grouping_key(strategy: Strategy) -> GenericStrategyKeyType: - raise NotImplementedError - - @staticmethod - @abc.abstractmethod - def get_metadata_type_for_group( - strategy_key: GenericStrategyKeyType, - ) -> Type[StrategyMetadata] | None: - raise NotImplementedError - - @staticmethod - @abc.abstractmethod - def sample_grouped_strategies( - group_key: GenericStrategyKeyType, - 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]: - """Sample grouped strategies. - - Returns: - - Sampled tokens - - Processed probs (whenever return_probs=True) - - Temperature (used to compute processed _log_ probs) - """ - 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 @@ -804,12 +714,11 @@ class BeamSearchSampleOnly(BeamSearchMixin, StrategyImplSampleOnly): ) -class FlashInferGroupedStrategySampler(GroupedStrategySampler[_STRATEGY_KEY_TYPE]): +class FlashInferGroupedStrategySampler: """Implements batched sampling with FlashInfer.sampling kernels.""" STRATEGY_KEY_TYPE: TypeAlias = _STRATEGY_KEY_TYPE - @override @staticmethod def strategy_grouping_key(strategy: Strategy) -> _STRATEGY_KEY_TYPE: match strategy: @@ -826,7 +735,6 @@ def strategy_grouping_key(strategy: Strategy) -> _STRATEGY_KEY_TYPE: case _: raise NotImplementedError("Unsupported strategy encountered") - @override @staticmethod def get_metadata_type_for_group( strategy_key: _STRATEGY_KEY_TYPE, @@ -837,7 +745,6 @@ def get_metadata_type_for_group( case _: return None - @override @staticmethod def sample_grouped_strategies( group_key: _STRATEGY_KEY_TYPE, @@ -849,6 +756,13 @@ def sample_grouped_strategies( return_probs: bool, group_metadata: StrategyMetadata | None = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: + """Sample grouped strategies. + + Returns: + - Sampled tokens + - Processed probs (whenever return_probs=True) + - Temperature (used to compute processed _log_ probs) + """ beam_width_in = 1 strategy_impl_cls: Type[_StrategyImpls.StrategyImpl] if return_probs: @@ -928,24 +842,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 +873,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 +891,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 diff --git a/tensorrt_llm/_torch/speculative/interface.py b/tensorrt_llm/_torch/speculative/interface.py index c601d8d4ca71..d0b0db6e3d32 100644 --- a/tensorrt_llm/_torch/speculative/interface.py +++ b/tensorrt_llm/_torch/speculative/interface.py @@ -55,9 +55,9 @@ def rejection_sampling_one_model( offset: Optional[int] = None, ) -> tuple[torch.Tensor, torch.Tensor]: # chain_speculative_sampling requires flashinfer>=0.6.4. This entry point can - # be reached independently of SpecWorkerBase.__init__'s use_flashinfer gate - # (e.g. via _can_use_rejection_sampling), so re-check the version here to fail - # with a clear message instead of a cryptic flashinfer error. + # be reached independently of SpecWorkerBase.__init__ (e.g. via + # _can_use_rejection_sampling), so re-check here to fail with a clear message + # instead of a cryptic flashinfer error. if not IS_FLASHINFER_AVAILABLE or Version( flashinfer.__version__) < Version("0.6.4"): raise RuntimeError( @@ -898,8 +898,15 @@ def __init__(self, use_separate_draft_kv_cache: bool = False): self.guided_decoder: Optional["CapturableGuidedDecoder"] = None self.force_num_accepted_tokens: float = get_force_num_accepted_tokens_float( ) - self.use_flashinfer = IS_FLASHINFER_AVAILABLE and Version( - flashinfer.__version__) >= Version("0.6.4") + # One-model speculative sampling goes through flashinfer unconditionally + # (sampling_batch_spec_dec_one_model), so flashinfer>=0.6.4 is a hard + # dependency here. Fail at construction with a clear error instead of + # crashing mid-inference on the first non-greedy sampling step. + if not IS_FLASHINFER_AVAILABLE or Version( + flashinfer.__version__) < Version("0.6.4"): + raise ImportError( + "Speculative decoding requires flashinfer>=0.6.4, please install " + "the version pinned in requirements.txt.") self.seed: Optional[torch.Tensor] = None self.offset: Optional[torch.Tensor] = None self.use_separate_draft_kv_cache = use_separate_draft_kv_cache @@ -1454,16 +1461,15 @@ def _draft_sampler_advanced( top_ks = spec_metadata.request_top_ks[:batch_size] top_ps = spec_metadata.request_top_ps[:batch_size] - if self.use_flashinfer: - if self.seed is None: - self.seed = torch.tensor([0], - dtype=torch.int64, - device=logits.device) - self.offset = torch.tensor([0], - dtype=torch.int64, - device=logits.device) - self.seed += 1 - self.seed %= (2**31) + if self.seed is None: + self.seed = torch.tensor([0], + dtype=torch.int64, + device=logits.device) + self.offset = torch.tensor([0], + dtype=torch.int64, + device=logits.device) + self.seed += 1 + self.seed %= (2**31) draft_tokens = sampling_batch_spec_dec_one_model(logits, temperatures, @@ -1687,17 +1693,16 @@ def _sample_tokens_for_batch( top_ks = spec_metadata.top_ks[:num_tokens] top_ps = spec_metadata.top_ps[:num_tokens] - if self.use_flashinfer: - # Lazily initialize seed/offset tensors on correct device - if self.seed is None: - self.seed = torch.tensor([0], - dtype=torch.int64, - device=logits.device) - self.offset = torch.tensor([0], - dtype=torch.int64, - device=logits.device) - self.seed += 1 - self.seed %= (2**31) + # Lazily initialize seed/offset tensors on correct device + if self.seed is None: + self.seed = torch.tensor([0], + dtype=torch.int64, + device=logits.device) + self.offset = torch.tensor([0], + dtype=torch.int64, + device=logits.device) + self.seed += 1 + self.seed %= (2**31) sampled_tokens = sampling_batch_spec_dec_one_model( logits, diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index fd3276163831..6a63ad9c8c49 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5166,12 +5166,6 @@ def validate_encoder_runtime_sizes(cls, v: Optional[int]) -> Optional[int]: # PrivateVars _quant_config: Optional[QuantConfig] = PrivateAttr(default=None) - disable_flashinfer_sampling: bool = Field( - default=False, - description= - "Disable the use of FlashInfer.sampling. This option is likely to be removed in the future.", - status="prototype") - max_stats_len: int = Field( default=1000, ge=-1, diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index db1c7772fdfb..023d4f07f4c1 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -307,13 +307,6 @@ "kind": "value", "path": "cuda_graph_config.seq_lens" }, - { - "allowed_values": [], - "annotation": "", - "converter": "", - "kind": "value", - "path": "disable_flashinfer_sampling" - }, { "allowed_values": [], "annotation": "", diff --git a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py index 5c8e3d62efd5..f55827f3e775 100644 --- a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py +++ b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py @@ -178,7 +178,6 @@ def _make_llm_args(): calibration_layer_indices=None, ), sampler_type=None, - disable_flashinfer_sampling=False, cuda_graph_config=None, parallel_config=SimpleNamespace(to_mapping=lambda: SimpleNamespace()), get_runtime_sizes=lambda: (1, 128, 128, 4), diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 95849080c370..9b4d1365a19e 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -56,9 +56,7 @@ def fixed_params(): return {"max_tokens": 8, "max_beam_width": 2} -@pytest.fixture(scope="module", - params=[("TRTLLMSampler", False), ("TorchSampler", False), - ("TorchSampler", True)]) +@pytest.fixture(scope="module", params=["TRTLLMSampler", "TorchSampler"]) def sampling_information(request): return request.param @@ -74,8 +72,7 @@ def model_kwargs(fixed_params, sampling_information) -> dict[str, Any]: weight_loader=DummyWeightLoader(), config_loader=DummyConfigLoader(), ), - sampler_type=sampling_information[0], - disable_flashinfer_sampling=sampling_information[1], + sampler_type=sampling_information, ) diff --git a/tests/unittest/_torch/sampler/test_logits_logprobs.py b/tests/unittest/_torch/sampler/test_logits_logprobs.py index 8fc1091bac94..ad07f70a3e5f 100644 --- a/tests/unittest/_torch/sampler/test_logits_logprobs.py +++ b/tests/unittest/_torch/sampler/test_logits_logprobs.py @@ -7,10 +7,7 @@ from utils.util import force_ampere from tensorrt_llm import LLM, SamplingParams -from tensorrt_llm._torch.pyexecutor.sampler.sampling_utils import ( - _StrategyImpls, - top_k_top_p_sampling_batch, -) +from tensorrt_llm._torch.pyexecutor.sampler.sampling_utils import _StrategyImpls from tensorrt_llm.executor.result import TokenLogprobs from tensorrt_llm.llmapi.llm_utils import KvCacheConfig @@ -93,13 +90,11 @@ def llm( yield llm -@pytest.fixture(scope="module", params=[False, True]) -def simple_llm(request) -> LLM: - disable_flashinfer_sampling = request.param +@pytest.fixture(scope="module") +def simple_llm() -> LLM: llm = LLM( model=os.path.join(llm_models_root(), "llama-models-v2", "TinyLlama-1.1B-Chat-v1.0"), max_batch_size=8, - disable_flashinfer_sampling=disable_flashinfer_sampling, kv_cache_config=global_kvcache_config_prompt_logprobs, ) return llm @@ -777,20 +772,15 @@ def test_processed_logprobs_e2e(logprobs_k: int, simple_llm: LLM): topp = topp if topp is not None else 1.0 temperature = temperature if temperature is not None else 1.0 - # perform maksing top-k top-p - if simple_llm.args.disable_flashinfer_sampling: - _, probs = top_k_top_p_sampling_batch( - logits_for_token, top_k=topk, top_p=topp, temperature=temperature - ) - else: - _, probs = _StrategyImpls.StrategyImplWithProbs._sample_with_probs( - logits_for_token, - group_logit_indices=None, - top_k=torch.tensor([topk], dtype=torch.int32, device="cuda"), - top_p=torch.tensor([topp], dtype=torch.float32, device="cuda"), - temperature=torch.tensor([temperature], dtype=torch.float32, device="cuda"), - generator=None, - ) + # perform masking top-k top-p via the flashinfer strategy impl + _, probs = _StrategyImpls.StrategyImplWithProbs._sample_with_probs( + logits_for_token, + group_logit_indices=None, + top_k=torch.tensor([topk], dtype=torch.int32, device="cuda"), + top_p=torch.tensor([topp], dtype=torch.float32, device="cuda"), + temperature=torch.tensor([temperature], dtype=torch.float32, device="cuda"), + generator=None, + ) if temperature != 0: logits_for_token /= temperature diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index 87e6b4b8dc01..5bda91081396 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -52,7 +52,6 @@ BeamSearch, FlashInferGroupedStrategySampler, Greedy, - SimpleGroupedStrategySampler, Strategy, StrategyMetadata, TemperatureOnly, @@ -1400,12 +1399,10 @@ def model_outputs( @pytest.fixture(scope="function") def sampler( self, - use_flashinfer: bool, max_draft_len: int, seq_slot_assignment: tuple[list[int], int], ) -> TorchSampler: return self._build_sampler( - use_flashinfer=use_flashinfer, max_draft_len=max_draft_len, seq_slot_assignment=seq_slot_assignment, ) @@ -1413,7 +1410,6 @@ def sampler( def _build_sampler( self, *, - use_flashinfer: bool, max_draft_len: int, seq_slot_assignment: tuple[list[int], int], ) -> TorchSampler: @@ -1425,7 +1421,6 @@ def _build_sampler( max_beam_width=1, # currently the only supported value max_num_sequences=num_seq_slots, max_total_draft_tokens=max_draft_len, - disable_flashinfer_sampling=(not use_flashinfer), disable_overlap_scheduler=False, ) ) @@ -1495,30 +1490,8 @@ def _mock_filter(self, requests: ScheduledRequests) -> list[LlmRequest]: new_tokens = new_tokens.squeeze(-1) return new_tokens - @pytest.mark.parametrize( - "use_flashinfer, max_draft_len, sampling_params_list", - [ - pytest.param(use_flashinfer, max_draft_len, []) - for (use_flashinfer, max_draft_len) in product( - [False, True], - [0, 3], - ) - ], - ) - def test_backend_selection( - self, - sampler: TorchSampler, - use_flashinfer: bool, - ): - """Check that TorchSampler uses the correct sampling backend.""" - expected_cls = ( - FlashInferGroupedStrategySampler if use_flashinfer else SimpleGroupedStrategySampler - ) - assert sampler._grouped_sampler_cls == expected_cls - @pytest.mark.parametrize( ( - "use_flashinfer", "max_draft_len", "draft_lens", "sampling_params_list", @@ -1529,18 +1502,16 @@ def test_backend_selection( [ # NB: non-zero draft len ensures that LlmRequest.py_target_probs is set. pytest.param( - use_flashinfer, 3, [3] * len(sampling_params_list), sampling_params_list, params_label, False, vocab_size, - id=f"{'FlashInfer' if use_flashinfer else 'Torch'}-{params_label}", + id=f"FlashInfer-{params_label}", ) # https://stackoverflow.com/a/75421799, does not work with nested loops - for (use_flashinfer, (sampling_params_list, params_label), vocab_size) in product( - [False, True], + for ((sampling_params_list, params_label), vocab_size) in product( _build_test_cases( vocab_size=VOCAB_SIZE, allow_greedy=False, # Greedy does not return probs @@ -1738,7 +1709,6 @@ def _uut(): def _compute_probs( self, *, - use_flashinfer: bool, model_outputs: dict[str, torch.Tensor], sampling_params_list: list[SamplingParams], seq_slot_assignment: tuple[list[int], int], @@ -1759,7 +1729,6 @@ def _compute_probs( # compute probs in general. draft_len_with_probs = max(1, max_draft_len) sampler_with_probs = self._build_sampler( - use_flashinfer=use_flashinfer, max_draft_len=draft_len_with_probs, seq_slot_assignment=seq_slot_assignment, ) @@ -1799,13 +1768,11 @@ def _inject_batching_check( patch_ctx: pytest.MonkeyPatch, *, sampler: TorchSampler, - use_flashinfer: bool, ): """Setup interception of sample_async and request grouping. - If FlashInfer.sampling is used, this validates that at every - invocation of sample_async, the sampling backend is called at most - once for any given sampling strategy (if FlashInfer.sampling is used). + Validates that at every invocation of sample_async, the FlashInfer + sampling backend is called at most once for any given sampling strategy. Used by test_samples. """ @@ -1814,76 +1781,76 @@ def _inject_batching_check( # This variable tracks which request types have been encountered. flashinfer_keys_seen: set[Any] = set() - if use_flashinfer: - assert sampler._grouped_sampler_cls == FlashInferGroupedStrategySampler - sample_grouped_strategies_orig = sampler._grouped_sampler_cls.sample_grouped_strategies - - def _sample_grouped_strategies( - group_key: FlashInferGroupedStrategySampler.STRATEGY_KEY_TYPE, - strategies: list[Strategy], - logits: torch.Tensor, - *, - group_logit_indices: Optional[torch.Tensor] = None, - generator: Optional[torch.Generator] = None, - return_probs: bool, - group_metadata: StrategyMetadata | None = None, - ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor] | float]: - assert generator is sampler.get_generator(logits.device) - if isinstance(group_key, tuple): - assert isinstance(group_key[0], str) - else: - assert isinstance(group_key, str) - nonlocal flashinfer_keys_seen - assert (group_key, return_probs) not in flashinfer_keys_seen - flashinfer_keys_seen.add((group_key, return_probs)) - result: tuple[ - torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor] | float - ] = sample_grouped_strategies_orig( + assert sampler._grouped_sampler_cls == FlashInferGroupedStrategySampler + sample_grouped_strategies_orig = sampler._grouped_sampler_cls.sample_grouped_strategies + + def _sample_grouped_strategies( + group_key: FlashInferGroupedStrategySampler.STRATEGY_KEY_TYPE, + strategies: list[Strategy], + logits: torch.Tensor, + *, + group_logit_indices: Optional[torch.Tensor] = None, + generator: Optional[torch.Generator] = None, + return_probs: bool, + group_metadata: StrategyMetadata | None = None, + ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor] | float]: + assert generator is sampler.get_generator(logits.device) + if isinstance(group_key, tuple): + assert isinstance(group_key[0], str) + else: + assert isinstance(group_key, str) + nonlocal flashinfer_keys_seen + assert (group_key, return_probs) not in flashinfer_keys_seen + flashinfer_keys_seen.add((group_key, return_probs)) + result: tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor] | float] = ( + sample_grouped_strategies_orig( group_key, strategies, logits, group_logit_indices=group_logit_indices, generator=generator, return_probs=return_probs, + group_metadata=group_metadata, ) - return result - - # _grouped_sampler_cls is a class; point the instance at a subclass - # that overrides the callable, rather than mutating the shared class. - instrumented_cls = type( - "InstrumentedFlashInferGroupedStrategySampler", - (FlashInferGroupedStrategySampler,), - {"sample_grouped_strategies": staticmethod(_sample_grouped_strategies)}, ) - patch_ctx.setattr(sampler, "_grouped_sampler_cls", instrumented_cls) + return result + + # _grouped_sampler_cls is a class; point the instance at a subclass + # that overrides the callable, rather than mutating the shared class. + instrumented_cls = type( + "InstrumentedFlashInferGroupedStrategySampler", + (FlashInferGroupedStrategySampler,), + {"sample_grouped_strategies": staticmethod(_sample_grouped_strategies)}, + ) + patch_ctx.setattr(sampler, "_grouped_sampler_cls", instrumented_cls) - sample_async_orig = sampler.sample_async + sample_async_orig = sampler.sample_async - def _sample_async( - scheduled_requests: ScheduledRequests, - model_outputs: dict[str, torch.Tensor], - num_context_logits_prefix_sum: list[int], - resource_manager=None, - ): - nonlocal flashinfer_keys_seen - flashinfer_keys_seen.clear() - res = sample_async_orig( - scheduled_requests, - model_outputs, - num_context_logits_prefix_sum, - resource_manager, - ) + def _sample_async( + scheduled_requests: ScheduledRequests, + model_outputs: dict[str, torch.Tensor], + num_context_logits_prefix_sum: list[int], + resource_manager=None, + ): + nonlocal flashinfer_keys_seen + flashinfer_keys_seen.clear() + res = sample_async_orig( + scheduled_requests, + model_outputs, + num_context_logits_prefix_sum, + resource_manager, + ) - # Fast greedy path bypasses flashinfer sampling, so flashinfer_keys_seen - # will be empty when all requests are greedy - all_greedy = all( - _request_strategy(req, vocab_size=2**31) == GREEDY - for req in scheduled_requests.all_requests() - ) - assert flashinfer_keys_seen or all_greedy - return res + # Fast greedy path bypasses flashinfer sampling, so flashinfer_keys_seen + # will be empty when all requests are greedy + all_greedy = all( + _request_strategy(req, vocab_size=2**31) == GREEDY + for req in scheduled_requests.all_requests() + ) + assert flashinfer_keys_seen or all_greedy + return res - patch_ctx.setattr(sampler, "sample_async", _sample_async) + patch_ctx.setattr(sampler, "sample_async", _sample_async) @dataclass(frozen=True, kw_only=True) class _TorchUtilsSamplingParams: @@ -2236,7 +2203,6 @@ def _validate_token_frequencies( @pytest.mark.parametrize( ( - "use_flashinfer", "max_draft_len", "sampling_params_list", "allow_zero_draft_len", @@ -2245,7 +2211,6 @@ def _validate_token_frequencies( ), [ pytest.param( - use_flashinfer, max_draft_len, sampling_params_list, allow_zero_draft_len, @@ -2256,21 +2221,19 @@ def _validate_token_frequencies( ), # bypass_sampling vocab_size, id=( - f"{'FlashInfer' if use_flashinfer else 'Torch'}" + f"FlashInfer" f"-draft_len={0 if allow_zero_draft_len else 1}..{max_draft_len}" f"-{params_label}" ), ) # https://stackoverflow.com/a/75421799, does not work with nested loops for ( - use_flashinfer, is_mixed, max_draft_len, allow_zero_draft_len, _build_test_cases, vocab_size, ) in product( - [False, True], [False, True], [0, 3], [False, True], @@ -2296,7 +2259,6 @@ def test_samples( sampling_params_list: list[SamplingParams], seq_slot_assignment: tuple[list[int], int], max_draft_len: int, - use_flashinfer: bool, allow_zero_draft_len: bool, # used by fixtures bypass_sampling: bool, monkeypatch: pytest.MonkeyPatch, @@ -2322,7 +2284,6 @@ def _uut_provider(is_warmup: bool) -> Generator[Callable[[], None], None, None]: # model_outputs. These probs, the computation of which is validated by 'test_probs', # are used to validate the batched sampling process later in this test. mock_requests_with_probs = self._compute_probs( - use_flashinfer=use_flashinfer, model_outputs=model_outputs, sampling_params_list=sampling_params_list, seq_slot_assignment=seq_slot_assignment, @@ -2338,9 +2299,7 @@ def _uut_provider(is_warmup: bool) -> Generator[Callable[[], None], None, None]: mock_sampling_log: Optional[list[TestBatchedSampling._MockSamplingLogEntry]] = None with monkeypatch.context() as patch_ctx: - self._inject_batching_check( - patch_ctx, sampler=sampler, use_flashinfer=use_flashinfer - ) + self._inject_batching_check(patch_ctx, sampler=sampler) if bypass_sampling: mock_sampling_log = self._instrument_sampling_backend( patch_ctx, sampler=sampler @@ -2502,7 +2461,6 @@ def _build_seq_slot_assignments() -> list[tuple[list[int], int, str]]: @pytest.mark.parametrize( ( - "use_flashinfer", "max_draft_len", "allow_zero_draft_len", "vocab_size", @@ -2511,7 +2469,6 @@ def _build_seq_slot_assignments() -> list[tuple[list[int], int, str]]: ), [ pytest.param( - False, # NB: _unbatch_sampling_results does not depend on backend max_draft_len, allow_zero_draft_len, vocab_size, @@ -2547,7 +2504,6 @@ def test_unbatch_sampling_results( vocab_size: int, # used by fixtures seq_slot_assignment: tuple[list[int], int], max_draft_len: int, - use_flashinfer: bool, # used by fixtures allow_zero_draft_len: bool, # used by fixtures ordered: bool, ): diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 1610f8700662..0db864d9e540 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -135,10 +135,6 @@ methods: annotation: bool default: False status: beta - disable_flashinfer_sampling: - annotation: bool - default: False - status: prototype moe_config: annotation: tensorrt_llm.llmapi.llm_args.MoeConfig status: beta