diff --git a/docs/source/features/sampling.md b/docs/source/features/sampling.md index d5a9eaef0685..72ad4a62ad41 100644 --- a/docs/source/features/sampling.md +++ b/docs/source/features/sampling.md @@ -75,7 +75,8 @@ llm.generate(["Hello, my name is", * The sampling is controlled via `SamplingParams`. -* By default (`temperature = top_p = top_k = None`), greedy sampling is used. +* By default (`temperature = top_p = top_k = None`), greedy sampling is used + (unless top-p decay is active, see below). * If either `temperature = 0`, `top_p = 0`, and/or `top_k = 1`, is specified, sampling is greedy, irrespective of the values of the remaining parameters. @@ -101,6 +102,19 @@ llm.generate(["Hello, my name is", * The implementation does not guarantee any particular treatment of tied probabilities. +* Top-P decay is supported: if `top_p_decay < 1` is specified, the effective `top_p` is + multiplied by `top_p_decay` after every sampled token, bounded from below by `top_p_min` + (default `1e-6`), and reset to the initial `top_p` whenever the token `top_p_reset_ids` + is sampled (default `-1`, which never matches a token). Out-of-range values + (`top_p_decay` or `top_p_min` outside `(0, 1]`, negative `top_p_reset_ids`) are rejected. + + * An active top-p decay implies top-p sampling even if `top_p` is unspecified or `top_p = 1` + (the initial `top_p` then defaults to 1). However, explicitly requested greedy sampling + (`temperature = 0`, `top_p = 0`, and/or `top_k = 1`) takes precedence over top-p decay. + + * Top-P decay is not supported in combination with beam search or with speculative decoding + modes that route draft tokens through the Torch Sampler; such requests are rejected. + ### Performance The Torch Sampler leverages the optimized sampling kernels provided by diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py index 5307876ae61d..8cfcca1e13b0 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py @@ -373,3 +373,105 @@ def gather_log_softmax(inputs_cuda: torch.Tensor, indices_cuda: torch.Tensor) -> torch._dynamo.mark_dynamic(inputs_cuda, 0) torch._dynamo.mark_dynamic(indices_cuda, 0) return Fusions._gather_log_softmax_impl(inputs_cuda, indices_cuda) + + # --- Top-P Decay ops --------------------------------------------------- + # Host-launch-bound per-step ops (a few dozen elements per row), fused with + # Inductor to keep the launch count low. mode="max-autotune-no-cudagraphs": + # cudagraphs is unsafe here (the update mutates persistent per-slot state + # in place and the gather's output is consumed outside the compiled region; + # cudagraph static output buffers get overwritten by subsequent replays). + # mark_dynamic on the batch-varying dims avoids recompilation as the batch + # composition changes. Compilation is lazy: the first decay-active request + # pays it (roughly a second); non-decay workloads never trigger it. See + # TorchSampler.TopPDecayStore for the feature-level semantics. + + @staticmethod + @torch.compile(mode="max-autotune-no-cudagraphs") + def _top_p_decay_update_impl( + runtime_top_p: torch.Tensor, + initial_top_p: torch.Tensor, + top_p_decay: torch.Tensor, + top_p_min: torch.Tensor, + reset_ids: torch.Tensor, + is_decay_slot: torch.Tensor, + step_tokens: torch.Tensor, + sampled_slots: torch.Tensor, + ) -> None: + active = is_decay_slot[sampled_slots] + current = runtime_top_p[sampled_slots] + updated = torch.where( + step_tokens[sampled_slots] == reset_ids[sampled_slots], + initial_top_p[sampled_slots], + torch.maximum(current * top_p_decay[sampled_slots], top_p_min[sampled_slots]), + ) + runtime_top_p[sampled_slots] = torch.where(active, updated, current) + + @staticmethod + def top_p_decay_update( + *, + runtime_top_p: torch.Tensor, + initial_top_p: torch.Tensor, + top_p_decay: torch.Tensor, + top_p_min: torch.Tensor, + reset_ids: torch.Tensor, + is_decay_slot: torch.Tensor, + step_tokens: torch.Tensor, + sampled_slots: torch.Tensor, + ) -> None: + """Fused in-place update of ``runtime_top_p`` for the sampled decay slots. + + Applies the Top-P Decay recurrence (see ``TorchSampler.TopPDecayStore`` + for the feature-level semantics) to every sampled row whose slot is + decay-active per ``is_decay_slot``. + + All per-slot tensors are 1-D of length ``max_num_sequences``; + ``step_tokens`` is a slot-indexed 1-D (possibly strided) int32 view of + the new-tokens buffer for a fixed step/beam + (``new_tokens[step, :, beam]``); ``sampled_slots`` is 1-D of length + ``num_sampled`` (this iteration's rows). ``runtime_top_p`` is mutated + in place; nothing is returned. + """ + torch._dynamo.mark_dynamic(sampled_slots, 0) + Fusions._top_p_decay_update_impl( + runtime_top_p, + initial_top_p, + top_p_decay, + top_p_min, + reset_ids, + is_decay_slot, + step_tokens, + sampled_slots, + ) + + @staticmethod + @torch.compile(mode="max-autotune-no-cudagraphs") + def _top_p_decay_gather_impl( + runtime_top_p: torch.Tensor, + is_decay_slot: torch.Tensor, + static_top_p: torch.Tensor, + slots: torch.Tensor, + ) -> torch.Tensor: + return torch.where(is_decay_slot[slots], runtime_top_p[slots], static_top_p) + + @staticmethod + def top_p_decay_gather( + *, + runtime_top_p: torch.Tensor, + is_decay_slot: torch.Tensor, + static_top_p: torch.Tensor, + slots: torch.Tensor, + ) -> torch.Tensor: + """Fused pre-sample per-row top-p gather for decay-active rows. + + Returns a new per-row tensor:: + + row_top_p[i] = runtime_top_p[slots[i]] if is_decay_slot[slots[i]] + = static_top_p[i] otherwise + + ``runtime_top_p`` / ``is_decay_slot`` are per-slot arrays; + ``static_top_p`` and ``slots`` are per-row (length = the group's + per-step row count). + """ + torch._dynamo.mark_dynamic(slots, 0) + torch._dynamo.mark_dynamic(static_top_p, 0) + return Fusions._top_p_decay_gather_impl(runtime_top_p, is_decay_slot, static_top_p, slots) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index bb50aab8c351..2f885dcbfffa 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -97,11 +97,13 @@ GenericStrategyKeyType, Strategy, StrategyMetadata, + TopPDecayMetadata, UtilsSamplingParams, get_rejected_indices, resolve_sampling_strategy, sample, sample_rejected, + top_p_decay_active, ) if sys.version_info[:2] >= (3, 12): @@ -475,9 +477,17 @@ def _get_max_beam_width(request: LlmRequest) -> int: def _request_get_sampling_params(request: LlmRequest) -> UtilsSamplingParams: sampling_config = request.sampling_config + # These sampling fields live on the C++ SamplingConfig as optional> + # (a shape designed for the batched TRT-LLM sampler); the torch sampler consumes + # them per request, so we unwrap the singleton lists into scalars here. When the + # TRT-LLM sampler is removed, this SamplingConfig-based plumbing should be removed + # too in favor of reading the values directly from the per-request params. temperature = _unwrap_singleton(cast(Optional[list[float]], sampling_config.temperature)) top_p = _unwrap_singleton(cast(Optional[list[float]], sampling_config.top_p)) top_k = _unwrap_singleton(cast(Optional[list[int]], sampling_config.top_k)) + top_p_decay = _unwrap_singleton(cast(Optional[list[float]], sampling_config.top_p_decay)) + top_p_min = _unwrap_singleton(cast(Optional[list[float]], sampling_config.top_p_min)) + top_p_reset_ids = _unwrap_singleton(cast(Optional[list[int]], sampling_config.top_p_reset_ids)) beam_width_out = _get_beam_width_out(request) beam_width_in = _get_beam_width_in(request) use_beam_search = _get_max_beam_width(request) > 1 @@ -489,6 +499,9 @@ def _request_get_sampling_params(request: LlmRequest) -> UtilsSamplingParams: beam_width_in=beam_width_in, beam_width_out=beam_width_out, use_beam_search=use_beam_search, + top_p_decay=top_p_decay, + top_p_min=top_p_min, + top_p_reset_ids=top_p_reset_ids, ) @@ -2247,6 +2260,79 @@ class LogProbsStore: """Shape: batch_size, max_tokens, max_topk_logprobs Usage: Stores the values of the topk logprobs""" + @dataclass(kw_only=True) + class TopPDecayStore: + """Per-slot runtime state for Top-P Decay -- the single source of truth + for the feature's semantics on the torch path. + + Semantics (matching the legacy C++ ``computeToppDecay`` kernel): after + every sampled token of a decay-active request:: + + runtime_top_p = initial_top_p if token == reset_id + = max(runtime_top_p * top_p_decay, top_p_min) otherwise + + A negative ``reset_ids`` sentinel (-1, "reset disabled") never matches, + since sampled token ids are non-negative. Decay is active iff + ``top_p_decay`` is set and < 1.0 (``SamplingParams. + params_imply_top_p_decay_active``); an active decay forces a + top-p-capable strategy even for an otherwise implicitly-greedy request + (initial top-p defaults to 1.0), while explicit greedy controls win. + Beam search and speculative draft tokens are rejected at admission + (``_validate_top_p_decay_request``); parameter ranges are enforced by + ``SamplingParams._validate`` / the executor::SamplingConfig constructor. + + Lifecycle per slot (each tensor has shape ``(max_num_sequences,)``): + + 1. Admission (``_setup_top_p_decay_for_new_requests``): membership is + cleared then re-set for the newly-admitted slots -- both the host-side + ``TorchSampler._top_p_decay_slots`` set (an O(1) hot-path early-out) + and its device mirror ``is_top_p_decay_slot_cuda`` (the gate the + fused ops use, so the hot path needs no host-side filtering) -- + and the per-slot buffers are initialized. This clear-then-set also + covers slot reuse: stale entries from a prior occupant are never + consumed. + 2. Pre-sample (``_build_top_p_decay_metadata`` -> ``TopPDecayMetadata`` + -> ``TopPDecayMixin``): the per-row top-p fed to top_p / + top_k_top_p sampling is overridden with the decayed runtime value + for decay-active rows (fused gather, ``top_p_decay_gather``). + 3. Post-sample (``_update_top_p_decay_after_sample``): the recurrence + above is applied in place for the sampled decay-active slots (fused + update, ``top_p_decay_update``). + 4. Finish (``_retire_top_p_decay_slot``): the slot leaves the + membership set so the early-outs re-arm; the device buffers need no + cleanup (a freed slot is never sampled, reuse re-initializes it). + """ + + runtime_top_p_decay_cuda: torch.Tensor + """The current (decaying) top-p per slot; mutated post-sample each step.""" + initial_top_p_decay_cuda: torch.Tensor + """The initial top-p per slot; used to reset on a reset-id match.""" + top_p_decay_cuda: torch.Tensor + """Per-slot multiplicative decay factor.""" + top_p_decay_min_cuda: torch.Tensor + """Per-slot lower bound for the decayed top-p.""" + top_p_decay_reset_ids_cuda: torch.Tensor + """Per-slot reset token id (< 0 never matches a sampled token).""" + is_top_p_decay_slot_cuda: torch.Tensor + """Per-slot bool gate (device mirror of ``_top_p_decay_slots``). Lets the + fused post-sample update op filter decay-active slots on the GPU, + avoiding a host-side ``.tolist()`` / set intersection each step.""" + + @classmethod + def create(cls, max_num_sequences: int) -> "TorchSampler.TopPDecayStore": + n = (max_num_sequences,) + return cls( + runtime_top_p_decay_cuda=torch.empty(n, dtype=torch.float32, device="cuda"), + initial_top_p_decay_cuda=torch.empty(n, dtype=torch.float32, device="cuda"), + top_p_decay_cuda=torch.empty(n, dtype=torch.float32, device="cuda"), + top_p_decay_min_cuda=torch.empty(n, dtype=torch.float32, device="cuda"), + top_p_decay_reset_ids_cuda=torch.empty(n, dtype=torch.int, device="cuda"), + # The gate buffer IS the gate, so (unlike the others) it must start + # False: a slot that was never admitted as decay-active must read + # False. + is_top_p_decay_slot_cuda=torch.zeros(n, dtype=torch.bool, device="cuda"), + ) + @dataclass(kw_only=True) class Store: new_tokens: torch.Tensor @@ -2258,6 +2344,8 @@ class Store: """Holds data related to beam search.""" log_probs_store: "TorchSampler.LogProbsStore" """Holds data related to log-probs handling.""" + top_p_decay_store: "TorchSampler.TopPDecayStore" + """Holds per-slot runtime state for Top-P Decay.""" def _create_store(self) -> Store: # Tensors necessary for all sampling methods @@ -2308,10 +2396,15 @@ def _create_store(self) -> Store: seq_offsets=seq_offsets, beam_idx_arange=beam_idx_arange, ) + # Per-slot Top-P Decay runtime state (FlashInfer path). Allocated for all + # sampler instances; only slots in self._top_p_decay_slots are ever read. + top_p_decay_store = self.TopPDecayStore.create(self.max_num_sequences) + return self.Store( new_tokens=new_tokens, log_probs_store=log_probs_store, beam_search_store=beam_search_store, + top_p_decay_store=top_p_decay_store, ) @dataclass(frozen=True, kw_only=True) @@ -2361,6 +2454,10 @@ def __init__(self, args: Args): "in requirements.txt." ) self._grouped_sampler_cls = FlashInferGroupedStrategySampler + # Slots with an active top-p-decay request. Sole gate for reading/updating + # the per-slot TopPDecayStore buffers; discarded on slot reuse so stale + # buffer entries are never consumed. + self._top_p_decay_slots: set[int] = set() # AutoDeploy build creates the sampler in inference mode, # which would disallow in-place mutating of new_tokens. @@ -2461,6 +2558,48 @@ def get_spec_tree_manager( def _use_beam_search(self) -> bool: return self.max_beam_width > 1 + def _validate_top_p_decay_request(self, request: LlmRequest) -> None: + """Reject unsupported combinations for a top-p-decay-active request. + + Top-p decay is supported only for single-token decode steps without beam + search. Called from validate_request (request admission), so a violating + request is failed individually instead of aborting the whole batch. + """ + params = _request_get_sampling_params(request) + # NB: value ranges need no re-check here. Every request enters through + # the executor::SamplingConfig constructor, which hard-validates + # top_p_decay in (0, 1], top_p_min in (0, 1] and top_p_reset_ids >= 0 + # (samplingConfig.cpp check* helpers) for all frontends. A reset id + # >= vocab_size is not checked anywhere but is semantically inert: it + # can never match a sampled token, i.e. it behaves as "reset disabled". + if not top_p_decay_active(params): + return + if params.use_beam_search: + raise ValueError("top_p_decay is not supported with beam search.") + # A non-zero draft length means the request carries speculative draft + # tokens and produces multiple tokens per step (req_num_steps = + # 1 + draft_token_length). One-model speculation (vanilla MTP, one-model + # Eagle3 / MTP-Eagle, SA, draft-target-one-model) uses its own + # SpecSamplerBase-derived sampler and never reaches TorchSampler; the + # drafter-based modes that DO flow draft tokens through TorchSampler + # (two-model draft-target, NGram, user-provided, two-model Eagle3 / + # MTP-Eagle) are what can make this length non-zero. top-p decay does not + # support these multi-token steps. + # NB: at admission time the draft tokens of drafter-based modes are + # usually not attached yet, so this check is best-effort. Two-model + # speculation (the only source of such requests in TorchSampler) is + # slated for removal, so in practice no speculative request reaches the + # decay path; a debug assert in _build_top_p_decay_metadata guards the + # invariant at sample time. + if get_draft_token_length(request) > 0: + raise ValueError( + "top_p_decay is not supported for requests carrying speculative " + "draft tokens (req_num_steps > 1). This covers the drafter-based " + "modes routed through TorchSampler (two-model draft-target, NGram, " + "user-provided, two-model Eagle3 / MTP-Eagle); one-model " + "speculation uses its own sampler and is unaffected." + ) + def _can_use_fast_greedy_path(self, requests: list[LlmRequest]) -> bool: """ Check if we can use the fast argmax path for greedy sampling. @@ -2921,6 +3060,10 @@ def _collect_new_requests_for_setup( @override def validate_request(self, request: LlmRequest) -> None: + # Reject unsupported top-p-decay combinations at admission, so only the + # offending request fails (raising later, inside setup_sampler_step or + # sampling, would abort the whole executor step). + self._validate_top_p_decay_request(request) if self._use_beam_search: if request.py_return_log_probs: if request.py_num_logprobs > 1: @@ -2995,6 +3138,10 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: all_sampling_requests=new_requests + scheduled_requests.generation_requests, ) + self._setup_top_p_decay_for_new_requests( + new_requests, new_seq_slots_cuda_long=seq_slots_tensor_cuda_long + ) + if self._use_beam_search: beam_search_store = self.store.beam_search_store assert beam_search_store is not None @@ -3005,6 +3152,106 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: max_prompt_len=max_prompt_len, ) + def _setup_top_p_decay_for_new_requests( + self, + new_requests: list[LlmRequest], + *, + new_seq_slots_cuda_long: torch.Tensor, + ) -> None: + """Refresh top-p-decay membership and per-slot buffers for admitted requests + (lifecycle step 1, see ``TopPDecayStore``). + + Drops stale membership from prior occupants of the newly-admitted slots + (host set and device gate), then re-admits the decay-active requests and + initializes their per-slot store entries. Unsupported decay combinations + were already rejected per-request in validate_request at admission. + """ + # Clear the device decay gate for every newly-admitted slot (covers slot + # reuse: a slot previously decay-active but reused by a non-decay request + # must read False). Decay-active slots are then set True below. + decay_gate = self.store.top_p_decay_store.is_top_p_decay_slot_cuda + decay_gate.index_fill_(0, new_seq_slots_cuda_long, False) + + decay_seq_slots: list[int] = [] + initial_top_p: list[float] = [] + top_p_decay: list[float] = [] + top_p_min: list[float] = [] + top_p_reset_ids: list[int] = [] + for request in new_requests: + slot = request.py_seq_slot + assert slot is not None + self._top_p_decay_slots.discard(slot) + sampling_params = _request_get_sampling_params(request) + if not top_p_decay_active(sampling_params): + continue + self._top_p_decay_slots.add(slot) + decay_seq_slots.append(slot) + # Initial runtime top-p defaults to 1.0 when top_p is unset. + initial_top_p.append( + sampling_params.top_p if sampling_params.top_p is not None else 1.0 + ) + # decay is guaranteed non-None and < 1.0 here (top_p_decay_active); + # min/reset fall back to the C++ runtime defaults when unset. + assert sampling_params.top_p_decay is not None + top_p_decay.append(sampling_params.top_p_decay) + top_p_min.append( + sampling_params.top_p_min if sampling_params.top_p_min is not None else 1e-6 + ) + top_p_reset_ids.append( + sampling_params.top_p_reset_ids + if sampling_params.top_p_reset_ids is not None + else -1 + ) + + if decay_seq_slots: + self._update_top_p_decay_store_for_new_requests( + decay_seq_slots=decay_seq_slots, + initial_top_p=initial_top_p, + top_p_decay=top_p_decay, + top_p_min=top_p_min, + top_p_reset_ids=top_p_reset_ids, + ) + + def _update_top_p_decay_store_for_new_requests( + self, + *, + decay_seq_slots: list[int], + initial_top_p: list[float], + top_p_decay: list[float], + top_p_min: list[float], + top_p_reset_ids: list[int], + ) -> None: + """Initialize per-slot Top-P Decay buffers for newly admitted decay requests. + + runtime_top_p and initial_top_p both start at the effective initial top-p; + the runtime value is decayed post-sample each step. + """ + store = self.store.top_p_decay_store + device = store.runtime_top_p_decay_cuda.device + slots_cuda = torch.tensor( + decay_seq_slots, device="cpu", dtype=torch.int64, pin_memory=prefer_pinned() + ).to(device, non_blocking=True) + floats_host = torch.tensor( + [initial_top_p, top_p_decay, top_p_min], + device="cpu", + dtype=torch.float32, + pin_memory=prefer_pinned(), + ) + floats_cuda = floats_host.to(device, non_blocking=True) + initial_cuda = floats_cuda[0] + reset_ids_cuda = torch.tensor( + top_p_reset_ids, device="cpu", dtype=torch.int32, pin_memory=prefer_pinned() + ).to(device, non_blocking=True) + + store.runtime_top_p_decay_cuda.index_copy_(0, slots_cuda, initial_cuda) + store.initial_top_p_decay_cuda.index_copy_(0, slots_cuda, initial_cuda) + store.top_p_decay_cuda.index_copy_(0, slots_cuda, floats_cuda[1]) + store.top_p_decay_min_cuda.index_copy_(0, slots_cuda, floats_cuda[2]) + store.top_p_decay_reset_ids_cuda.index_copy_(0, slots_cuda, reset_ids_cuda) + # Enable the device gate for these decay-active slots (cleared for all new + # slots in setup_sampler_step just before this call). + store.is_top_p_decay_slot_cuda.index_fill_(0, slots_cuda, True) + @staticmethod def _prepare_beam_search( beam_search_store: BeamSearchStore, @@ -3490,6 +3737,7 @@ def _add_metadata_to_grouped_requests( *, seq_slots_cuda: torch.Tensor, seq_lens_cuda: torch.Tensor, + req_num_steps: torch.Tensor, ) -> dict[RequestGroupKey[GenericStrategyKeyType], RequestGroupValueWithMetadata]: grouped_requests_with_metadata: dict[ RequestGroupKey[GenericStrategyKeyType], RequestGroupValueWithMetadata @@ -3499,6 +3747,7 @@ def _add_metadata_to_grouped_requests( num_requests = len(requests) for key, value in grouped_requests.items(): metadata_type = get_metadata_type_for_group_fn(key.strategy_key) + metadata: StrategyMetadata | None if metadata_type is BeamSearchMetadata: assert beam_search_store is not None assert seq_lens is not None, "seq_lens is required for beam search" @@ -3527,6 +3776,13 @@ def _add_metadata_to_grouped_requests( seq_offsets=beam_search_store.seq_offsets, beam_idx_arange=beam_search_store.beam_idx_arange, ) + elif metadata_type is TopPDecayMetadata: + metadata = self._build_top_p_decay_metadata( + group_req_indices=value.indices, + req_num_steps=req_num_steps, + seq_slots=seq_slots, + seq_slots_cuda=seq_slots_cuda, + ) elif metadata_type is None: metadata = None else: @@ -3682,6 +3938,7 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None: for req_idx, req in enumerate(state.requests): if req.state == LlmRequestState.GENERATION_COMPLETE: + self._retire_top_p_decay_slot(req) continue if req.py_beam_width > 1: @@ -3748,6 +4005,20 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None: self._finish_reasons_handler.store.num_accepted_draft_tokens_host[ req.py_seq_slot ] = req.py_num_accepted_draft_tokens + if req.state == LlmRequestState.GENERATION_COMPLETE: + self._retire_top_p_decay_slot(req) + + def _retire_top_p_decay_slot(self, req: LlmRequest) -> None: + """Retire a finished request's slot from the top-p-decay membership set + (lifecycle step 4, see ``TopPDecayStore``), so the O(1) hot-path + early-outs re-arm once decay traffic drains. + + Callers ensure ``req`` has finished. Requests that finish outside the + sampler (e.g. cancellation) are covered by the slot-reuse cleanup at + admission instead. + """ + if self._top_p_decay_slots and req.py_seq_slot is not None: + self._top_p_decay_slots.discard(req.py_seq_slot) def _return_log_probs(self, requests: list[LlmRequest]) -> bool: return any(req.py_return_log_probs for req in requests) @@ -4015,6 +4286,65 @@ def provision_bias_index() -> int: # sharing). logits[logits_bias_mask_cuda] += biases_tensor_cuda + def _build_top_p_decay_metadata( + self, + *, + group_req_indices: torch.Tensor, + req_num_steps: torch.Tensor, + seq_slots: torch.Tensor, + seq_slots_cuda: torch.Tensor, + ) -> Optional[TopPDecayMetadata]: + """Build the Top-P Decay metadata for a top_p / top_k_top_p group. + + Lifecycle step 2, see ``TopPDecayStore``. Returns None when no request + currently uses decay. The metadata's ``slots`` tensor is aligned to the + group's per-STEP row order (matching group_strategies_per_step); + non-decay rows (possibly multi-step draft rows) are gated out on-device + by ``is_decay_slot``, so decay presence in the group is not checked + host-side: a group without decay rows samples every row with its static + top-p -- same result as returning None. + """ + if not self._top_p_decay_slots: + return None + store = self.store.top_p_decay_store + # Fast path (steady-state decoding): if every row in the group is + # single-token, the per-STEP row order equals the per-request order, and + # (group_req_indices being sorted ascending) a contiguous group's slots + # are just a slice of seq_slots_cuda -- no host layout build and no H2D + # copy. + first_req = int(group_req_indices[0].item()) + last_req = int(group_req_indices[-1].item()) + group_steps = req_num_steps[group_req_indices] + if last_req - first_req + 1 == group_req_indices.size(0) and ( + group_steps.max().item() == 1 + ): + per_step_slots_cuda = seq_slots_cuda[first_req : last_req + 1] + else: + # Build the per-STEP slot layout (each request contributes + # req_num_steps rows). + group_seq_slots = seq_slots[group_req_indices] + if __debug__: + # Internal invariant (stripped under python -O): a decay-active + # row is always single-token -- the only source of multi-step + # rows in TorchSampler is two-model speculation (slated for + # removal), and decay + draft tokens is rejected per-request at + # admission in validate_request. + decay_row_steps = group_steps[ + torch.isin(group_seq_slots, torch.tensor(list(self._top_p_decay_slots))) + ] + assert decay_row_steps.numel() == 0 or decay_row_steps.max().item() == 1, ( + "top_p_decay row with req_num_steps != 1; decay + draft tokens " + "should have been rejected at admission" + ) + per_step_slots_cuda = torch.repeat_interleave(group_seq_slots.long(), group_steps).to( + seq_slots_cuda.device, non_blocking=True + ) + return TopPDecayMetadata( + slots=per_step_slots_cuda, + runtime_top_p=store.runtime_top_p_decay_cuda, + is_decay_slot=store.is_top_p_decay_slot_cuda, + ) + @nvtx_range("sample_batched_by_strategy") @torch.inference_mode() def _sample_batched_by_strategy( @@ -4051,6 +4381,7 @@ def _sample_batched_by_strategy( get_metadata_type_for_group_fn=self._grouped_sampler_cls.get_metadata_type_for_group, seq_slots_cuda=seq_slots_cuda, seq_lens_cuda=seq_lens_cuda, + req_num_steps=req_num_steps, ) generator_cuda = self.get_generator(cuda_device) @@ -4284,6 +4615,7 @@ def _unbatch_sampling_results( new_tokens_cuda: torch.Tensor, req_num_generated_tokens: torch.Tensor, seq_slots: torch.Tensor, + seq_slots_cuda: torch.Tensor, ) -> torch.Tensor: batch_req_indices = batched_sampling_result.batch_req_indices batch_next_tokens_cuda_int = batched_sampling_result.batch_next_tokens_cuda_int @@ -4317,8 +4649,48 @@ def _dims_canonically_ordered(t: torch.Tensor) -> bool: new_tokens_cuda.view(-1, *new_tokens_cuda.shape[2:]).scatter_( 0, batch_dest_indices_1d_cuda, batch_next_tokens_cuda_int ) + # Post-sample: decay the runtime top-p for any decay-active slots that were + # sampled this iteration (must run after tokens land in new_tokens_cuda). + # batch_req_indices is a permutation of all sampled requests, so the set of + # sampled slots is exactly seq_slots (the kernel updates each slot + # independently; order is irrelevant) -- pass the resident device copy + # instead of gathering seq_slots[batch_req_indices] on host and copying it. + self._update_top_p_decay_after_sample( + new_tokens_cuda=new_tokens_cuda, + sampled_slots_cuda=seq_slots_cuda, + ) return self._copy_to_host(new_tokens_cuda) + def _update_top_p_decay_after_sample( + self, + *, + new_tokens_cuda: torch.Tensor, + sampled_slots_cuda: torch.Tensor, + ) -> None: + """Apply the post-sample decay recurrence for sampled decay-active slots. + + See ``TopPDecayStore`` for the feature-level semantics (lifecycle + step 3). Restricting to the sampled slots avoids reading stale + new_tokens_cuda for slots that were not scheduled this iteration. + """ + # Host-side O(1) early-out: skip the kernel launch when no request uses + # decay; otherwise a single fused (torch.compile) op gates on + # is_decay_slot on-device and gathers the sampled token in place (decay is + # single-token-only, so the token is at local step 0, beam 0). + if not self._top_p_decay_slots: + return + store = self.store.top_p_decay_store + Fusions.top_p_decay_update( + runtime_top_p=store.runtime_top_p_decay_cuda, + initial_top_p=store.initial_top_p_decay_cuda, + top_p_decay=store.top_p_decay_cuda, + top_p_min=store.top_p_decay_min_cuda, + reset_ids=store.top_p_decay_reset_ids_cuda, + is_decay_slot=store.is_top_p_decay_slot_cuda, + step_tokens=new_tokens_cuda[DEFAULT_STEP_IDX, :, DEFAULT_BEAM_IDX], + sampled_slots=sampled_slots_cuda, + ) + @staticmethod @torch.inference_mode() def _apply_min_length_penalty( @@ -4953,6 +5325,7 @@ def _process_requests( new_tokens_cuda=new_tokens_cuda, req_num_generated_tokens=sampling_requests_metadata.req_num_generated_tokens, seq_slots=seq_slots_host, + seq_slots_cuda=seq_slots_cuda, ) # NB: update_requests syncs w/ device computation and async D2H copies diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py index f99946ffd5a1..7ca58c7edc08 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py @@ -106,6 +106,10 @@ class UtilsSamplingParams: use_beam_search: Whether to use beam search. beam_width_in: The beam_width of a request before the sampling step. beam_width_out: The beam_width of a request after the sampling step. + top_p_decay: Per-step multiplicative decay applied to the runtime top-p. + top_p_min: Lower bound for the decayed runtime top-p. + top_p_reset_ids: Token id which, when sampled, resets the runtime top-p to + its initial value. A value < 0 never matches a token. """ temperature: Optional[float] @@ -114,6 +118,38 @@ class UtilsSamplingParams: use_beam_search: Optional[bool] beam_width_in: Optional[int] = None beam_width_out: Optional[int] = None + top_p_decay: Optional[float] = None + top_p_min: Optional[float] = None + top_p_reset_ids: Optional[int] = None + + +@dataclass(kw_only=True) +class TopPDecayMetadata(StrategyMetadata): + """Per-group runtime top-p override for Top-P Decay (attached to top_p / + top_k_top_p groups via the ``StrategyMetadata`` mechanism). + + ``slots`` maps each per-step group row to its sequence slot; the decayed + per-row top-p is gathered on-device from the per-slot ``runtime_top_p`` + store, gated by ``is_decay_slot`` (non-decay rows keep their static top-p). + Consumed by the TopP*/TopKTopP* strategy impls in ``sample()``. See + ``TorchSampler.TopPDecayStore`` for the feature-level semantics. + """ + + slots: torch.Tensor + """Per-step group rows' sequence slots (int64, device).""" + runtime_top_p: torch.Tensor + """Per-slot runtime (decayed) top-p store (float32, device).""" + is_decay_slot: torch.Tensor + """Per-slot decay-active gate (bool, device).""" + + +def top_p_decay_active(params: UtilsSamplingParams) -> bool: + """Whether dynamic top-p decay is active for a request. + + Delegates to the single-source predicate on SamplingParams; note that + ``top_p_min`` / ``top_p_reset_ids`` alone do not activate dynamic behavior. + """ + return SamplingParams.params_imply_top_p_decay_active(params.top_p_decay) def resolve_sampling_strategy(params: UtilsSamplingParams, *, vocab_size: int) -> Strategy: @@ -124,11 +160,15 @@ def resolve_sampling_strategy(params: UtilsSamplingParams, *, vocab_size: int) - top_p = params.top_p top_k = params.top_k + # The greedy verdict (including the top-p-decay override of the implicit + # all-unset greedy default, and explicit greedy controls winning over decay) + # is single-sourced in SamplingParams.params_imply_greedy_decoding. if SamplingParams.params_imply_greedy_decoding( temperature=temperature, top_p=top_p, top_k=top_k, use_beam_search=use_beam_search, + top_p_decay=params.top_p_decay, ): return GREEDY @@ -152,7 +192,10 @@ def resolve_sampling_strategy(params: UtilsSamplingParams, *, vocab_size: int) - assert top_k > 1, "non-greedy sampling requires valid top_k" need_top_k = top_k < vocab_size assert top_p > 0, "non-greedy sampling requires valid top_p" - need_top_p = top_p < 1 + # A decay-active request must go through a top-p-capable path even when its + # initial top_p is 1.0, so the runtime top-p (sourced per-row at sample time) + # can shrink the nucleus on later steps. + need_top_p = top_p < 1 or top_p_decay_active(params) if need_top_p: if need_top_k: @@ -346,6 +389,34 @@ def _sample_with_probs( new_tokens = cls._sample_from_probs(probs, generator=generator) return new_tokens, probs + class TopPDecayMixin: + """Mixed into the TopP*/TopKTopP* impls (the owners of a per-row + ``_top_p`` tensor) to consume ``TopPDecayMetadata``.""" + + _top_p: torch.Tensor + + def _maybe_apply_top_p_decay(self, group_metadata: Optional[StrategyMetadata]) -> None: + """Override the per-row static top-p with the decayed runtime top-p. + + Only decay-active rows (per the on-device ``is_decay_slot`` gate) are + overridden, so a group mixing top-p-decay and plain top-p requests + keeps each row's correct value. The overridden ``self._top_p`` tensor + then feeds both sampling and ``top_p_renorm_probs_op`` (so processed + logprobs match). Fused via torch.compile (gather + gate + select). + """ + if not isinstance(group_metadata, TopPDecayMetadata): + return + assert self._top_p.shape == group_metadata.slots.shape, ( + self._top_p.shape, + group_metadata.slots.shape, + ) + self._top_p = Fusions.top_p_decay_gather( + runtime_top_p=group_metadata.runtime_top_p, + is_decay_slot=group_metadata.is_decay_slot, + static_top_p=self._top_p, + slots=group_metadata.slots, + ) + class StrategyImplWithProbs(StrategyImpl): @override @classmethod @@ -374,7 +445,7 @@ def sample( ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: return self._sample_greedy_with_probs(logits, group_logit_indices=group_logit_indices) - class TopKTopPWithProbs(StrategyImplWithProbs): + class TopKTopPWithProbs(TopPDecayMixin, StrategyImplWithProbs): def __init__(self, top_k: torch.Tensor, top_p: torch.Tensor, temperature: torch.Tensor): self._top_k = top_k self._top_p = top_p @@ -400,6 +471,7 @@ def sample( generator: Optional[torch.Generator] = None, group_metadata: Optional[StrategyMetadata] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + self._maybe_apply_top_p_decay(group_metadata) return self._sample_with_probs( logits, group_logit_indices=group_logit_indices, @@ -442,7 +514,7 @@ def sample( generator=generator, ) - class TopPWithProbs(StrategyImplWithProbs): + class TopPWithProbs(TopPDecayMixin, StrategyImplWithProbs): def __init__(self, top_p: torch.Tensor, temperature: torch.Tensor): self._top_p = top_p self._temperature = temperature @@ -466,6 +538,7 @@ def sample( generator: Optional[torch.Generator] = None, group_metadata: Optional[StrategyMetadata] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + self._maybe_apply_top_p_decay(group_metadata) return self._sample_with_probs( logits, group_logit_indices=group_logit_indices, @@ -534,7 +607,7 @@ def sample( logits = logits[group_logit_indices] return torch.argmax(logits, dim=-1), None - class TopKTopPSampleOnly(StrategyImplSampleOnly): + class TopKTopPSampleOnly(TopPDecayMixin, StrategyImplSampleOnly): def __init__(self, top_k: torch.Tensor, top_p: torch.Tensor, temperature: torch.Tensor): self._top_k = top_k self._top_p = top_p @@ -560,6 +633,7 @@ def sample( generator: Optional[torch.Generator] = None, group_metadata: Optional[StrategyMetadata] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + self._maybe_apply_top_p_decay(group_metadata) logits = self._prepare_logits_with_temperature( logits, group_logit_indices, self._temperature ) @@ -605,7 +679,7 @@ def sample( check_nan=self._flashinfer_check_nans(probs), ), None - class TopPSampleOnly(StrategyImplSampleOnly): + class TopPSampleOnly(TopPDecayMixin, StrategyImplSampleOnly): def __init__(self, top_p: torch.Tensor, temperature: torch.Tensor): self._top_p = top_p self._temperature = temperature @@ -629,6 +703,7 @@ def sample( generator: Optional[torch.Generator] = None, group_metadata: Optional[StrategyMetadata] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + self._maybe_apply_top_p_decay(group_metadata) probs = self._prepare_probs_with_temperature( logits, group_logit_indices, self._temperature ) @@ -755,6 +830,8 @@ def get_metadata_type_for_group( match strategy_key: case ("beam_search", _, _): return BeamSearchMetadata + case "top_p" | "top_k_top_p": + return TopPDecayMetadata case _: return None diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index bded9477cd3d..e7fed5e2ced0 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -210,9 +210,9 @@ class SamplingParams: If neither temperature, top_p, nor top_k are specified, sampling is greedy. If temperature > 0 and/or top_k > 1 are specified, sampling will proceed accordingly and top_p will default to top_p = 1. Setting top_p = 0 should result in greedy sampling, but is currently disallowed in the backend. - top_p_min (float, optional): Controls decay in the top-P algorithm. topPMin is lower-bound. None means using C++ runtime default 1.e-6. Defaults to None. - top_p_reset_ids (int, optional): Controls decay in the top-P algorithm. Indicates where to reset the decay. None means using C++ runtime default 1. Defaults to None. - top_p_decay (float, optional): Controls decay in the top-P algorithm. The decay value. None means using C++ runtime default 1.f. Defaults to None. + top_p_min (float, optional): Controls decay in the top-P algorithm. topPMin is lower-bound. Must be in (0, 1]; invalid values are rejected. None means using C++ runtime default 1.e-6. Defaults to None. + top_p_reset_ids (int, optional): Controls decay in the top-P algorithm. The token id which, when sampled, resets the decayed top-P to its initial value. Must be >= 0; invalid values are rejected. None means using C++ runtime default -1 (which never matches a token). Defaults to None. + top_p_decay (float, optional): Controls decay in the top-P algorithm. The decay value. Must be in (0, 1]; invalid values are rejected. None means using C++ runtime default 1.f. Defaults to None. seed (int, optional): Controls the random seed used by the random number generator in sampling. None means using C++ runtime default 0. Defaults to None. temperature (float, optional): Controls the modulation of logits when sampling new tokens. It can have values >= 0.f. Defaults to None. The value None is treated as "not specified" in the following. @@ -379,6 +379,20 @@ def _validate(self): if self.temperature is not None and self.temperature < 0: raise ValueError(f"require temperature >= 0, got temperature={self.temperature}") + # Top-p decay param ranges mirror the hard checks in the + # executor::SamplingConfig constructor (samplingConfig.cpp check* + # helpers); rejecting here gives a clear, early error instead of a + # RuntimeError from the C++ boundary. Note top_p_min > top_p is + # intentionally allowed (the runtime top-p may rise toward top_p_min). + if self.top_p_decay is not None and not 0.0 < self.top_p_decay <= 1.0: + raise ValueError(f"require 0 < top_p_decay <= 1, got top_p_decay={self.top_p_decay}") + if self.top_p_min is not None and not 0.0 < self.top_p_min <= 1.0: + raise ValueError(f"require 0 < top_p_min <= 1, got top_p_min={self.top_p_min}") + if self.top_p_reset_ids is not None and self.top_p_reset_ids < 0: + raise ValueError( + f"require top_p_reset_ids >= 0, got top_p_reset_ids={self.top_p_reset_ids}" + ) + if self.best_of is not None and self.best_of < self.n: raise ValueError(f"best_of ({self.best_of}) cannot be less than n ({self.n})") @@ -427,8 +441,37 @@ def _validate(self): if self.logprobs_simple_format and self.use_beam_search: raise ValueError("logprobs_simple_format is not supported with beam search") - # NB: Static, because downstream code only holds instances of - # bindings.SamplingConfig (not SamplingParams). + # NB: The predicates below are static because downstream code (e.g. + # sampling_utils.resolve_sampling_strategy) only holds instances of + # bindings.SamplingConfig (not SamplingParams). They are the single + # source of truth for the greedy / top-p-decay resolution shared by + # _greedy_decoding and the torch sampler. + + @staticmethod + def params_imply_top_p_decay_active(top_p_decay: Optional[float]) -> bool: + """Whether dynamic top-p decay is active. + + Active iff ``top_p_decay`` is explicitly set and ``< 1.0``; a decay of + ``1.0`` (the C++ default) is a no-op. Values outside ``(0, 1]`` are + rejected up front (_validate and the executor::SamplingConfig + constructor), so they never reach this predicate. + """ + return top_p_decay is not None and top_p_decay < 1.0 + + @staticmethod + def params_imply_explicit_greedy( + *, + temperature: Optional[float], + top_p: Optional[float], + top_k: Optional[int], + ) -> bool: + """Whether the request carries an explicit greedy control. + + Explicit means top_k == 1, top_p == 0.0, or temperature == 0, as opposed + to the implicit "all params unset" greedy default. + """ + return top_k == 1 or top_p == 0.0 or temperature == 0 + @staticmethod def params_imply_greedy_decoding( *, @@ -436,13 +479,23 @@ def params_imply_greedy_decoding( top_p: Optional[float], top_k: Optional[int], use_beam_search: bool | None, - ): - return (not use_beam_search) and ( - (temperature is None and top_p is None and top_k is None) - or top_k == 1 - or top_p == 0.0 - or temperature == 0 - ) + top_p_decay: Optional[float] = None, + ) -> bool: + """Whether the parameters resolve to greedy decoding. + + An explicit greedy control always wins. The implicit "all params unset" + greedy default is overridden by an active top-p decay (which implies + top-p sampling so the decayed runtime top-p can take effect); callers + that do not support decay may omit ``top_p_decay``. + """ + if use_beam_search: + return False + if SamplingParams.params_imply_explicit_greedy( + temperature=temperature, top_p=top_p, top_k=top_k + ): + return True + implicitly_greedy = temperature is None and top_p is None and top_k is None + return implicitly_greedy and not SamplingParams.params_imply_top_p_decay_active(top_p_decay) @property def _greedy_decoding(self) -> bool: @@ -451,6 +504,7 @@ def _greedy_decoding(self) -> bool: top_p=self.top_p, top_k=self.top_k, use_beam_search=self.use_beam_search, + top_p_decay=self.top_p_decay, ) @property diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index e2e0d1bc9402..892b93357fb0 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -59,6 +59,7 @@ TopKTopP, TopP, UtilsSamplingParams, + resolve_sampling_strategy, ) from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests from tensorrt_llm.bindings import SamplingConfig @@ -2577,6 +2578,11 @@ class UutResultWrapper: result: Optional[UutResult] = None res = UutResultWrapper() + # Precomputed outside the no-sync region (mirrors the production + # resident device copy of seq_slots). + seq_slots_tensor_cuda = ( + seq_slots_tensor.to(torch.int64).pin_memory().to("cuda", non_blocking=True) + ) def _uut(res=res): new_tokens_host = sampler._unbatch_sampling_results( @@ -2584,6 +2590,7 @@ def _uut(res=res): new_tokens_cuda=new_tokens_cuda, req_num_generated_tokens=req_num_steps, seq_slots=seq_slots_tensor, + seq_slots_cuda=seq_slots_tensor_cuda, ) res.result = UutResult(new_tokens_host=new_tokens_host) @@ -2756,3 +2763,133 @@ def test_stale_conditional_and_unconditional_same_cell(self): logits = self._run_stale([req], [True], {0: 9}) assert self._banned_cols(logits[0]) == {2} assert not torch.isnan(logits).any() + + +class TestTopPDecay: + """Minimal functional guards for Top-P Decay in TorchSampler. + + Covers strategy routing, the post-sample runtime update (parity with the + C++ computeToppDecay recurrence; cases ported from + topPSamplingLayerTest.cpp), and per-request rejection of unsupported + combinations. + """ + + VOCAB_SIZE = 1000 + + @staticmethod + def _params(**kw) -> UtilsSamplingParams: + base = dict(temperature=None, top_p=None, top_k=None, use_beam_search=False) + base.update(kw) + return UtilsSamplingParams(**base) + + @staticmethod + def _make_sampler(*, max_draft_len=0): + return TorchSampler( + TorchSampler.Args( + max_seq_len=128, + max_draft_len=max_draft_len, + max_num_sequences=8, + max_beam_width=1, + max_total_draft_tokens=max_draft_len, + disable_overlap_scheduler=True, + ) + ) + + def test_strategy_routing(self): + # Active decay (set and < 1.0) forces a top-p-capable strategy even for + # an otherwise-greedy request (initial top-p defaults to 1.0), so the + # decayed runtime value can take effect on later steps. + s = resolve_sampling_strategy(self._params(top_p_decay=0.5), vocab_size=self.VOCAB_SIZE) + assert s[0] == "top_p" and s[1] == pytest.approx(1.0) + s = resolve_sampling_strategy( + self._params(top_k=50, top_p=0.9, top_p_decay=0.8), vocab_size=self.VOCAB_SIZE + ) + assert s[0] == "top_k_top_p" + # decay == 1.0 (the C++ default) is a no-op and does not activate... + s = resolve_sampling_strategy(self._params(top_p_decay=1.0), vocab_size=self.VOCAB_SIZE) + assert s is GREEDY + # ...and an explicit greedy control wins over an active decay. + s = resolve_sampling_strategy( + self._params(top_p_decay=0.5, top_k=1), vocab_size=self.VOCAB_SIZE + ) + assert s is GREEDY + + def test_runtime_update_parity(self): + # Post-sample update parity with the C++ computeToppDecay recurrence + # (a negative reset_id never matches, since token ids are non-negative): + # runtime = initial if token == reset_id + # = max(runtime * decay, min) otherwise + sampler = self._make_sampler() + store = sampler.store.top_p_decay_store + configs = [ + dict(initial=0.8, decay=0.3, top_p_min=0.5, reset_id=2), # decay, then reset + dict(initial=0.2, decay=0.9, top_p_min=0.1, reset_id=-1), # plain decay, floored + dict(initial=0.3, decay=0.5, top_p_min=0.6, reset_id=-1), # min > initial: rises + ] + token_steps = [[1, 2, 3], [9, 9, 9], [9, 9, 9]] + slots = list(range(len(configs))) + for slot, cfg in zip(slots, configs): + sampler._top_p_decay_slots.add(slot) + store.runtime_top_p_decay_cuda[slot] = cfg["initial"] + store.initial_top_p_decay_cuda[slot] = cfg["initial"] + store.top_p_decay_cuda[slot] = cfg["decay"] + store.top_p_decay_min_cuda[slot] = cfg["top_p_min"] + store.top_p_decay_reset_ids_cuda[slot] = cfg["reset_id"] + store.is_top_p_decay_slot_cuda[slot] = True + + runtime = [cfg["initial"] for cfg in configs] + slots_cuda = torch.tensor(slots, dtype=torch.int64, device="cuda") + for step in range(3): + for slot in slots: + sampler.store.new_tokens[0, slot, 0] = token_steps[slot][step] + sampler._update_top_p_decay_after_sample( + new_tokens_cuda=sampler.store.new_tokens, sampled_slots_cuda=slots_cuda + ) + got = store.runtime_top_p_decay_cuda.cpu() + for slot, cfg in zip(slots, configs): + tok = token_steps[slot][step] + if tok == cfg["reset_id"]: + runtime[slot] = cfg["initial"] + else: + runtime[slot] = max(runtime[slot] * cfg["decay"], cfg["top_p_min"]) + assert got[slot].item() == pytest.approx(runtime[slot], abs=1e-6), (step, slot) + + @staticmethod + def _mock_request(params: SamplingParams, *, draft_tokens=None): + params._validate() + req = SimpleNamespace( + sampling_config=SamplingConfig(params._get_sampling_config()), + is_context_init_state=False, + py_sampling_strategy=None, + py_draft_tokens=draft_tokens, + ) + req.get_beam_width_by_iter = lambda for_next_iteration=False: 1 + return cast(LlmRequest, req) + + @pytest.mark.parametrize( + "bad_kwargs", + [ + {"top_p_decay": 1.5}, + {"top_p_decay": -0.5}, + {"top_p_decay": 0.0}, + {"top_p_min": 0.0}, + {"top_p_min": 1.5}, + {"top_p_reset_ids": -1}, + ], + ) + def test_out_of_range_decay_params_rejected(self, bad_kwargs): + # Out-of-range decay params raise (mirroring the executor::SamplingConfig + # constructor's hard checks) instead of the former warn-and-default. + with pytest.raises(ValueError): + SamplingParams(**bad_kwargs) + + def test_reject_speculative_draft_tokens(self): + # Decay + draft tokens through TorchSampler is rejected per-request at + # admission (validate_request), so only the offending request fails. + sampler = self._make_sampler(max_draft_len=4) + with pytest.raises(ValueError, match="speculative"): + sampler.validate_request( + self._mock_request(SamplingParams(top_p=0.9, top_p_decay=0.5), draft_tokens=[1, 2]) + ) + # Same request without decay is accepted. + sampler.validate_request(self._mock_request(SamplingParams(top_p=0.9), draft_tokens=[1, 2]))