diff --git a/docs/source/features/sampling.md b/docs/source/features/sampling.md index 29f75abde6ba..ea741b622e4b 100644 --- a/docs/source/features/sampling.md +++ b/docs/source/features/sampling.md @@ -13,6 +13,7 @@ The PyTorch backend supports a wide variety of features, listed below: | | Rejection sampling (composable) | Return Logits | | | | Return LogProbs | | | | TopK LogProbs | +| | | Penalties | ## General usage @@ -115,6 +116,33 @@ llm.generate(["Hello, my name is", * 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. +* Occurrence penalties are supported: `repetition_penalty`, `presence_penalty` and + `frequency_penalty` discourage (or encourage) the model from reusing tokens it has + already seen. All three rewrite the logits before temperature scaling, driven by the + occurrence history of the prompt plus everything generated so far. Writing `c` for the + number of times a token has occurred in that history: + + * `repetition_penalty` (default `1.0`) rescales the logit of every token with `c > 0`: + the logit is divided by the penalty when it is non-negative and multiplied by it when + it is negative. The two branches move a positive and a negative logit the same way, so + a value `> 1` always pushes a seen token down, and a value `< 1` always pulls it up. + Must be `> 0`. + + * `presence_penalty` (default `0.0`) subtracts the penalty itself from every token with + `c > 0`. The amount does not depend on `c`, so it controls whether a token reappears, + not how often. + + * `frequency_penalty` (default `0.0`) subtracts the penalty multiplied by `c`, so the + more often a token has already been produced, the harder it is pushed down. + + * `prompt_ignore_length` (default `0`) excludes the first N prompt tokens from the + presence and frequency counts. Those ignored tokens still count for + `repetition_penalty`. Values `<= 0` have no effect, and values larger than the prompt + are clamped to the prompt length. + + * Occurrence penalties are not supported in combination with beam search; such requests + are rejected. + * If `no_repeat_ngram_size = n` is specified, any token that would recreate an `n`-gram already present in the sequence (prompt included) is excluded from sampling. `None` or `0` disables the restriction. diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py index e62568c76df5..d2918cb9ac9e 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py @@ -456,3 +456,178 @@ def top_p_decay_gather( 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) + + # --- Occurrence penalties (repetition / presence / frequency) ----------- + # torch/torch.compile counterpart of the C++ ``batchApplyPenalty`` kernel, + # driven by ``PenaltyHandler`` in penalties.py, which owns the + # workspace and documents its semantics (see ``PenaltyStore`` there). + + @staticmethod + def update_occurrence_workspace( + counts_cuda: torch.Tensor, + presence_prefix_cuda: Optional[torch.Tensor], + counted_slots: torch.Tensor, + counted_tokens: torch.Tensor, + prefix_slots: Optional[torch.Tensor] = None, + prefix_tokens: Optional[torch.Tensor] = None, + ) -> None: + """Scatter (slot, token) pairs into the persistent occurrence workspace. + + Args: + counts_cuda: ``int32[num_slots, vocab_size]``, incremented in place. + presence_prefix_cuda: ``bool[num_slots, vocab_size]`` prefix-presence + mask, or ``None`` when no active request uses + ``prompt_ignore_length``. + counted_slots / counted_tokens: 1-D int64 pairs to increment in + ``counts_cuda``. + prefix_slots / prefix_tokens: 1-D int64 pairs to mark in + ``presence_prefix_cuda``; ``None`` when there is nothing to mark. + """ + if counted_slots.numel() > 0: + ones = torch.ones( + counted_slots.shape[0], dtype=counts_cuda.dtype, device=counts_cuda.device + ) + # accumulate=True sums repeated (slot, token) pairs -> occurrence count. + counts_cuda.index_put_((counted_slots, counted_tokens), ones, accumulate=True) + if ( + presence_prefix_cuda is not None + and prefix_slots is not None + and prefix_tokens is not None + and prefix_slots.numel() > 0 + ): + # Marking a dense bool mask is idempotent, so duplicate tokens are safe. + presence_prefix_cuda[prefix_slots, prefix_tokens] = True + + # fullgraph=True is safe here: served model has fixed shapes and compiles ~2 graphs, + # well under the default limit (8) + @staticmethod + @torch.compile(fullgraph=True, mode="max-autotune-no-cudagraphs") + def _apply_occurrence_penalties_impl( + logits: torch.Tensor, + counts_cuda: torch.Tensor, + prefix_seen_cuda: Optional[torch.Tensor], + active_cuda: torch.Tensor, + has_previous_token_cuda: torch.Tensor, + new_tokens: torch.Tensor, + seq_slots: torch.Tensor, + request_offsets: torch.Tensor, + request_num_steps: torch.Tensor, + repetition_cuda: torch.Tensor, + presence_cuda: torch.Tensor, + frequency_cuda: torch.Tensor, + ) -> None: + vocab = logits.size(-1) + + # Fold the device-pending sampled token into the persistent counts, once per armed + # active slot, before the gather reads them, via one flat scatter_add. Masked entries + # add 0 at counts[slot, 0], so inactive/unarmed/out-of-range slots are no-ops. + previous_token = new_tokens[0, seq_slots, 0].to(torch.int64) + fold_ok = ( + active_cuda[seq_slots] + & has_previous_token_cuda[seq_slots] + & (request_num_steps > 0) + & (previous_token >= 0) + & (previous_token < vocab) + ) + flat_index = seq_slots * vocab + torch.where( + fold_ok, previous_token, previous_token.new_zeros(()) + ) + counts_cuda.view(-1).scatter_add_(0, flat_index, fold_ok.to(counts_cuda.dtype)) + + # Map each logits row to its owning request with a broadcasted range comparison. + # This is O(T * R), but T and R are both small (rows per step x requests) and the + # whole thing fuses into the surrounding elementwise graph, so it measures faster + # than either a searchsorted lookup or a repeat_interleave expansion. Notably + # torch.repeat_interleave must NOT be used here: its output length is + # sum(num_steps), which -- without an explicit host-provided output_size -- torch + # reads back from the device, and that per-step D2H sync destroys the overlap + # between the sampler's host work and the model forward (measured ~20x slower). + rows = torch.arange(logits.size(0), device=logits.device).unsqueeze(1) # [T, 1] + owned = (rows >= request_offsets) & (rows < request_offsets + request_num_steps) # [T, R] + row_owned = owned.any(dim=1) # [T] + row_slot = (owned * seq_slots).sum(dim=1) # [T]; slot per row, 0 for unowned + row_active = row_owned & active_cuda[row_slot] + + count = counts_cuda[row_slot] + rep = repetition_cuda[row_slot].unsqueeze(1) + pre = presence_cuda[row_slot].unsqueeze(1) + freq = frequency_cuda[row_slot].unsqueeze(1) + + seen = count > 0 + if prefix_seen_cuda is not None: + # Prompt-ignore-prefix tokens count for repetition only, not presence/frequency. + seen = seen | prefix_seen_cuda[row_slot] + + penalized = logits.float() + repeated = torch.where(penalized < 0, penalized * rep, penalized / rep) + penalized = torch.where(seen, repeated, penalized) + penalized = penalized - torch.where( + count > 0, + pre + freq * count.to(torch.float32), + penalized.new_zeros(()), + ) + limit = torch.finfo(logits.dtype).max + penalized = penalized.clamp(-limit, limit).to(logits.dtype) + # Cast before the select so inactive rows stay bit-identical, then write in place. + logits.copy_(torch.where(row_active.unsqueeze(1), penalized, logits)) + + @staticmethod + def apply_batched_occurrence_penalties( + logits: torch.Tensor, + counts_cuda: torch.Tensor, + presence_prefix_cuda: Optional[torch.Tensor], + active_cuda: torch.Tensor, + has_previous_token_cuda: torch.Tensor, + new_tokens: torch.Tensor, + seq_slots: torch.Tensor, + request_offsets: torch.Tensor, + request_num_steps: torch.Tensor, + repetition_cuda: torch.Tensor, + presence_cuda: torch.Tensor, + frequency_cuda: torch.Tensor, + ) -> None: + """Apply occurrence penalties to ``logits`` in place, before temperature handling. + + Args: + logits: ``[T, vocab_size]`` packed generated-token logits, where + ``T == sum(num_steps * num_beams)``. Request ``r`` owns the rows + ``request_offsets[r] + step`` for ``step in [0, request_num_steps[r])``; + rows no request owns are left bit-identical. Modified in place. + counts_cuda / presence_prefix_cuda: the occurrence workspace; see + ``PenaltyHandler.PenaltyStore`` for their semantics. + active_cuda / has_previous_token_cuda / repetition_cuda / presence_cuda / + frequency_cuda: per-slot buffers of length ``max_num_sequences``. + new_tokens: ``[max_tokens, max_num_sequences, max_beam_width]`` device + buffer holding the previous step's sampled token. + seq_slots: ``int64[R]`` slot per request. + request_offsets / request_num_steps: ``[R]`` device tensors, already + staged by the caller. The owned spans must not overlap, but they need + not be ordered, and rows they skip are left bit-identical. + + All heavy lifting is fused into the single compiled ``_apply_occurrence_penalties_impl`` + graph; this wrapper only marks the batch-varying dims dynamic. + """ + if seq_slots.numel() == 0 or logits.size(0) == 0: + return + + # Batch-varying dim-0 tensors; mark every one, or an unmarked peer forces the marked + # dims to specialize (ConstraintViolationError under dynamic=None). counts/active/params + # keep dim 0 == max_num_sequences (fixed) and new_tokens dim 1 == max_num_sequences. + torch._dynamo.mark_dynamic(logits, 0) + torch._dynamo.mark_dynamic(seq_slots, 0) + torch._dynamo.mark_dynamic(request_offsets, 0) + torch._dynamo.mark_dynamic(request_num_steps, 0) + Fusions._apply_occurrence_penalties_impl( + logits, + counts_cuda, + presence_prefix_cuda, + active_cuda, + has_previous_token_cuda, + new_tokens, + seq_slots, + request_offsets, + request_num_steps, + repetition_cuda, + presence_cuda, + frequency_cuda, + ) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/penalties.py b/tensorrt_llm/_torch/pyexecutor/sampler/penalties.py new file mode 100644 index 000000000000..4f3ea18a2613 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/sampler/penalties.py @@ -0,0 +1,483 @@ +# 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. + +"""Occurrence penalties (repetition / presence / frequency) for ``TorchSampler``. + +The feature's persistent device state lives in :class:`PenaltyStore` (which +documents the workspace semantics) and its whole lifecycle in +:class:`PenaltyHandler`; ``TorchSampler`` owns one instance and drives it +through request validation, admission, the per-step apply, and the +post-processing commit of finalized tokens. +""" + +from dataclasses import dataclass + +import torch + +from tensorrt_llm._utils import nvtx_range, prefer_pinned + +from ..llm_request import LlmRequest +from .ops.vanilla import Fusions +from .sampler_common import _get_max_beam_width, _unwrap_singleton + +__all__ = ["PenaltyHandler", "PenaltyStore"] + + +def _has_occurrence_penalty(request: LlmRequest) -> bool: + sampling_config = request.sampling_config + repetition = _unwrap_singleton(sampling_config.repetition_penalty) + presence = _unwrap_singleton(sampling_config.presence_penalty) + frequency = _unwrap_singleton(sampling_config.frequency_penalty) + return ( + (repetition is not None and repetition != 1.0) + or (presence is not None and presence != 0.0) + or (frequency is not None and frequency != 0.0) + ) + + +@dataclass(kw_only=True) +class PenaltyStore: + """Persistent device state: penalty-parameter buffers + occurrence workspace. + + This is the torch counterpart of the tensors ``PenaltyLayer`` allocates, and + the anchor for the workspace semantics the ops and the handler rely on: + + * The **parameter buffers** (``repetition_cuda`` / ``presence_cuda`` / + ``frequency_cuda``, plus the ``active_cuda`` gate) are the counterpart of + ``allocateBuffer`` + ``fillBuffers``: one entry per sequence slot, written + once per request and gathered every step, never rebuilt on the host. + * The **occurrence workspace** (``counts_cuda`` and ``presence_prefix_cuda``) + is the counterpart of ``allocateWorkspace`` / ``mPenaltyWorkspaceDevice``, + updated incrementally each step. A token in the ignored prompt prefix + ``[0, prompt_ignore_length)`` only sets ``presence_prefix_cuda``, so it + contributes to the repetition penalty but not to presence/frequency; every + other token (the rest of the prompt plus each generated token) increments + ``counts_cuda``, which drives presence/frequency and -- via ``counts > 0`` -- + repetition as well. + """ + + max_num_sequences: int + device: torch.device + + # --- Penalty parameters (allocateBuffer counterpart), shape [max_num_sequences] --- + repetition_cuda: torch.Tensor + """float32; per-slot repetition penalty (default 1.0).""" + presence_cuda: torch.Tensor + """float32; per-slot presence penalty (default 0.0).""" + frequency_cuda: torch.Tensor + """float32; per-slot frequency penalty (default 0.0).""" + active_cuda: torch.Tensor + """bool[slots]; whether a slot has an active occurrence penalty.""" + has_previous_token_cuda: torch.Tensor + """bool[slots]; whether ``new_tokens`` contains a token to accumulate.""" + + # --- Occurrence workspace (allocateWorkspace counterpart), allocated lazily --- + counts_cuda: torch.Tensor | None = None + """int32[slots, vocab_size] or None; occurrence counts (see class docstring).""" + presence_prefix_cuda: torch.Tensor | None = None + """bool[slots, vocab_size] or None; ignored-prompt-prefix presence mask.""" + + # Per-step request metadata, staged into persistent device buffers by + # ``stage_request_metadata`` so the hot path does not allocate per step. + request_offsets_cuda: torch.Tensor | None = None + request_num_steps_cuda: torch.Tensor | None = None + + @classmethod + def create(cls, *, max_num_sequences: int, device: torch.device) -> "PenaltyStore": + """Allocate the vocab-independent buffers with their no-op defaults. + + ``inference_mode(False)`` guards every allocation in this class: the + buffers persist across sampler steps and are mutated in place later, which + inference-mode tensors forbid. + """ + with torch.inference_mode(False): + return cls( + max_num_sequences=max_num_sequences, + device=device, + repetition_cuda=torch.ones(max_num_sequences, dtype=torch.float32, device=device), + presence_cuda=torch.zeros(max_num_sequences, dtype=torch.float32, device=device), + frequency_cuda=torch.zeros(max_num_sequences, dtype=torch.float32, device=device), + active_cuda=torch.zeros(max_num_sequences, dtype=torch.bool, device=device), + has_previous_token_cuda=torch.zeros( + max_num_sequences, dtype=torch.bool, device=device + ), + ) + + def ensure_workspace(self, *, vocab_size: int, needs_prefix: bool) -> None: + """Allocate the vocab-sized workspace on first use. + + Deferred because ``vocab_size`` is only known once logits arrive, mirroring + ``PenaltyLayer::allocateWorkspace`` being gated on penalty usage. The prefix + mask is allocated only if some request has used ``prompt_ignore_length``. + """ + with torch.inference_mode(False): + if self.counts_cuda is None: + self.counts_cuda = torch.zeros( + (self.max_num_sequences, vocab_size), + dtype=torch.int32, + device=self.device, + ) + if needs_prefix and self.presence_prefix_cuda is None: + self.presence_prefix_cuda = torch.zeros( + (self.max_num_sequences, vocab_size), + dtype=torch.bool, + device=self.device, + ) + + def stage_request_metadata( + self, request_offsets_host: torch.Tensor, request_num_steps_host: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Copy this step's ``[R]`` request metadata into persistent device buffers. + + The host tensors are already pinned by the caller, so each step costs two + small async H2D copies into a reused allocation rather than two fresh + device tensors. Returned views are only valid until the next call. + """ + num_requests = request_offsets_host.numel() + with torch.inference_mode(False): + if ( + self.request_offsets_cuda is None + or self.request_offsets_cuda.numel() < num_requests + ): + capacity = max(num_requests, self.max_num_sequences) + self.request_offsets_cuda = torch.empty( + capacity, dtype=request_offsets_host.dtype, device=self.device + ) + self.request_num_steps_cuda = torch.empty( + capacity, dtype=request_num_steps_host.dtype, device=self.device + ) + assert self.request_num_steps_cuda is not None + offsets = self.request_offsets_cuda[:num_requests] + num_steps = self.request_num_steps_cuda[:num_requests] + offsets.copy_(request_offsets_host, non_blocking=True) + num_steps.copy_(request_num_steps_host, non_blocking=True) + return offsets, num_steps + + +class PenaltyHandler: + """Applies the occurrence penalties: repetition, presence and frequency. + + These rescale or subtract from a token's logit based on how often it has already + occurred, and run before the sampling strategy divides by temperature. Bans that + force a logit to -inf (min_length, bad words, no-repeat-ngram) are a different + kind of transform and live in ``TokenBanHandler``. + + The implementation follows the C++ ``batchApplyPenalty`` kernel + (``cpp/tensorrt_llm/kernels/penaltyKernels.cu``) as driven by ``PenaltyLayer``. + Its persistent device state lives in :class:`PenaltyStore`, which documents the + workspace semantics. Per-slot parameter buffers are filled once per request, + batched across all requests admitted in a step (``prepare_for_new_request`` + accumulates on the host, ``update_for_new_requests`` issues the device updates). + Vocab-sized workspaces are allocated lazily and skipped entirely when no matching + request uses an occurrence penalty. + """ + + @dataclass(kw_only=True) + class _SlotState: + """Per-slot host-only bookkeeping (never read by the ops).""" + + prompt_ignore_length: int + initialized: bool = False + + def __init__( + self, + *, + max_num_sequences: int, + device: torch.device | str, + ): + self._max_num_sequences = max_num_sequences + self._device = torch.device(device) + # Whether any (past or current) active request uses prompt_ignore_length > 0, + # which requires allocating the presence-prefix mask. + self._needs_prefix = False + self._num_active_slots = 0 + # Per-slot state; None marks a slot without active occurrence penalties. + self._slots: list[PenaltyHandler._SlotState | None] = [None] * max_num_sequences + # Slots admitted this step that carry an occurrence penalty, with their + # parameters; drained by ``update_for_new_requests``. + self._new_slots: list[int] = [] + self._new_repetition: list[float] = [] + self._new_presence: list[float] = [] + self._new_frequency: list[float] = [] + self.store = PenaltyStore.create(max_num_sequences=max_num_sequences, device=self._device) + + @staticmethod + def validate_request(request: LlmRequest) -> None: + """Reject unsupported combinations for a penalized request. + + Called from ``TorchSampler.validate_request`` (request admission), so a + violating request is failed individually instead of aborting the whole batch. + """ + if _get_max_beam_width(request) > 1 and _has_occurrence_penalty(request): + raise ValueError( + "TorchSampler does not support repetition, presence, or frequency " + "penalties with beam search." + ) + + def _to_device(self, values: list[int], dtype: torch.dtype) -> torch.Tensor: + return torch.tensor(values, dtype=dtype, pin_memory=prefer_pinned()).to( + self._device, non_blocking=True + ) + + def prepare_for_new_request(self, request: LlmRequest, slot: int) -> None: + """Record the slot's penalty parameters for this step's batched flush. + + Called from ``TorchSampler.setup_sampler_step`` for each new request, mirroring + ``PenaltyLayer::setup`` (``fillBuffers`` + per-``batchSlot`` ``setZero``). This + only touches host state; ``update_for_new_requests`` issues the device updates + for all requests admitted in the step at once. Inactive slots are never + gathered, so their stale parameters/counts are left untouched. + """ + was_active = self._slots[slot] is not None + if not (_get_max_beam_width(request) == 1 and _has_occurrence_penalty(request)): + self._slots[slot] = None + if was_active: + self._num_active_slots -= 1 + return + + sampling_config = request.sampling_config + repetition = _unwrap_singleton(sampling_config.repetition_penalty) + presence = _unwrap_singleton(sampling_config.presence_penalty) + frequency = _unwrap_singleton(sampling_config.frequency_penalty) + prompt_ignore_length = _unwrap_singleton(sampling_config.prompt_ignore_length) + # min(prompt_ignore_length, inputLen), matching the C++ kernel. + prompt_ignore_length = min( + prompt_ignore_length if prompt_ignore_length is not None else 0, + request.py_orig_prompt_len, + ) + if prompt_ignore_length > 0: + self._needs_prefix = True + + self._slots[slot] = self._SlotState(prompt_ignore_length=prompt_ignore_length) + if not was_active: + self._num_active_slots += 1 + + self._new_slots.append(slot) + self._new_repetition.append(repetition if repetition is not None else 1.0) + self._new_presence.append(presence if presence is not None else 0.0) + self._new_frequency.append(frequency if frequency is not None else 0.0) + + def update_for_new_requests(self, *, new_seq_slots_cuda_long: torch.Tensor) -> None: + """Flush this step's admissions to the device in a handful of batched updates. + + ``new_seq_slots_cuda_long`` holds *every* slot admitted this step. Clearing the + active gate and the pending-token flag across all of them also covers slot + reuse: a slot whose prior occupant was penalized but whose new occupant is not + must read False. + """ + store = self.store + store.active_cuda.index_fill_(0, new_seq_slots_cuda_long, False) + store.has_previous_token_cuda.index_fill_(0, new_seq_slots_cuda_long, False) + + if not self._new_slots: + return + + slots_cuda = self._to_device(self._new_slots, torch.int64) + # One [3, N] host tensor -> one H2D for all three parameter buffers. + params_cuda = torch.tensor( + [self._new_repetition, self._new_presence, self._new_frequency], + dtype=torch.float32, + pin_memory=prefer_pinned(), + ).to(self._device, non_blocking=True) + store.repetition_cuda.index_copy_(0, slots_cuda, params_cuda[0]) + store.presence_cuda.index_copy_(0, slots_cuda, params_cuda[1]) + store.frequency_cuda.index_copy_(0, slots_cuda, params_cuda[2]) + store.active_cuda.index_fill_(0, slots_cuda, True) + + # Re-zero the workspace rows so a prior occupant's counts do not leak in. + if store.counts_cuda is not None: + store.counts_cuda.index_fill_(0, slots_cuda, 0) + if store.presence_prefix_cuda is not None: + store.presence_prefix_cuda.index_fill_(0, slots_cuda, False) + + self._new_slots.clear() + self._new_repetition.clear() + self._new_presence.clear() + self._new_frequency.clear() + + def _initialize_workspace( + self, + request: LlmRequest, + state: "PenaltyHandler._SlotState", + vocab_size: int, + ) -> None: + """Initialize one regular slot from its prompt exactly once.""" + if state.initialized: + return + + slot = request.py_seq_slot + assert slot is not None + counts_cuda = self.store.counts_cuda + assert counts_cuda is not None + + prompt = request.get_tokens(0)[: request.py_orig_prompt_len] + state.initialized = True + if not prompt: + return + + # One conversion for the whole prompt; the split point is just + # prompt_ignore_length, so the two groups are plain slices. + tokens = self._to_device(prompt, torch.int64) + prefix_tokens = tokens[: state.prompt_ignore_length] + counted_tokens = tokens[state.prompt_ignore_length :] + + # Multimodal models place placeholder ids >= vocab_size in the prompt (see + # _torch/models/modeling_multimodal_utils.py), so out-of-range ids reach us + # here and must be dropped before they index the workspace. + counted_tokens = counted_tokens[(counted_tokens >= 0) & (counted_tokens < vocab_size)] + prefix_tokens = prefix_tokens[(prefix_tokens >= 0) & (prefix_tokens < vocab_size)] + + Fusions.update_occurrence_workspace( + counts_cuda, + self.store.presence_prefix_cuda, + torch.full_like(counted_tokens, slot), + counted_tokens, + torch.full_like(prefix_tokens, slot), + prefix_tokens, + ) + + def update_token_counts( + self, + updates: list[tuple[int, list[int]]], + ) -> None: + """Commit finalized sampled tokens that replaced the device pending token. + + This is used after sampler-side postprocessing has finalized a multi-token + result. The complete confirmed sequence is counted here, then the raw first + token left in ``new_tokens`` is marked consumed so the next kernel cannot count + it again. Regular one-token sampling never calls this method and keeps its + fused device-pending fast path. + """ + if not updates or self._num_active_slots == 0: + return + + counts_cuda = self.store.counts_cuda + assert counts_cuda is not None + vocab_size = counts_cuda.size(-1) + consumed_slots: list[int] = [] + counted_slots: list[int] = [] + counted_tokens: list[int] = [] + + for slot, tokens in updates: + if self._slots[slot] is None: + continue + consumed_slots.append(slot) + for token in tokens: + if 0 <= token < vocab_size: + counted_slots.append(slot) + counted_tokens.append(token) + + if consumed_slots: + self.store.has_previous_token_cuda.index_fill_( + 0, self._to_device(consumed_slots, torch.int64), False + ) + + if not counted_tokens: + return + + Fusions.update_occurrence_workspace( + counts_cuda, + self.store.presence_prefix_cuda, + self._to_device(counted_slots, torch.int64), + self._to_device(counted_tokens, torch.int64), + ) + + @nvtx_range("apply_penalties") + @torch.inference_mode() + def apply( + self, + logits: torch.Tensor, + requests: list[LlmRequest], + *, + new_tokens: torch.Tensor, + seq_slots: torch.Tensor, + request_offsets: torch.Tensor, + request_num_steps: torch.Tensor, + is_draft_batch: bool = False, + ) -> None: + """Apply the occurrence penalties to ``logits`` in place. + + ``logits`` is the packed generated-token logits ``[sum(num_steps * num_beams), + vocab_size]``; request ``r`` owns ``request_num_steps[r]`` consecutive rows + starting at ``request_offsets[r]``, in beam-major / step-minor order. + ``request_offsets`` / ``request_num_steps`` are the caller's pinned host + tensors and are staged to the device here. + + Args: + is_draft_batch: draft batches share this sampler but draw ``py_seq_slot`` + from a separate numbering space that collides with target slots, so + penalizing them would read/write an unrelated target request's + occurrence state; skip them like the pending-steps tracking. + """ + if is_draft_batch or not requests or self._num_active_slots == 0: + return + + # Cheap per-batch scan so the vocab-sized workspace is only allocated when this + # batch actually contains a penalized request. + active_requests: list[tuple[LlmRequest, "PenaltyHandler._SlotState"]] = [] + for request in requests: + slot = request.py_seq_slot + assert slot is not None + state = self._slots[slot] + if state is not None: + active_requests.append((request, state)) + if not active_requests: + return + + store = self.store + store.ensure_workspace(vocab_size=logits.size(-1), needs_prefix=self._needs_prefix) + counts_cuda = store.counts_cuda + assert counts_cuda is not None + for request, state in active_requests: + self._initialize_workspace(request, state, logits.size(-1)) + + request_offsets_cuda, request_num_steps_cuda = store.stage_request_metadata( + request_offsets, request_num_steps + ) + Fusions.apply_batched_occurrence_penalties( + logits, + counts_cuda, + store.presence_prefix_cuda, + store.active_cuda, + store.has_previous_token_cuda, + new_tokens, + seq_slots, + request_offsets_cuda, + request_num_steps_cuda, + store.repetition_cuda, + store.presence_cuda, + store.frequency_cuda, + ) + # Arm has_previous_token for the slots this call penalized (active, num_steps > 0) + # so the next apply folds their sampled new_tokens. Done here rather than in the + # compiled op because the op's fold reads the flag for every request row; flipping + # it in the same graph would make the result depend on execution order within the + # kernel. + # + # The scan is kept on the host deliberately. The same thing can be expressed on + # device as active_cuda[seq_slots] & (num_steps > 0), avoiding this loop and the + # H2D, but that costs several extra kernel launches and measured 5-7us slower for + # batches up to 32 and no better at 64-256: the loop overlaps with the model + # forward, the launches do not. + pending_token_slots: list[int] = [] + for request, num_steps in zip(requests, request_num_steps.tolist()): + slot = request.py_seq_slot + if slot is None: + continue + if self._slots[slot] is not None and num_steps > 0: + pending_token_slots.append(slot) + if pending_token_slots: + store.has_previous_token_cuda.index_fill_( + 0, self._to_device(pending_token_slots, torch.int64), True + ) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 42d7c68b273d..db8546ff67f2 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -98,6 +98,7 @@ get_logprobs_from_request, store_logprobs_list_to_request, ) +from .penalties import PenaltyHandler from .sampler_common import ( DEFAULT_BEAM_IDX, DEFAULT_STEP_IDX, @@ -1424,6 +1425,10 @@ def __init__(self, args: Args): self.store.new_tokens.shape == self._finish_reasons_handler.store.finish_reasons_cuda.shape ) + self._penalty_handler = PenaltyHandler( + max_num_sequences=self.max_num_sequences, + device="cuda", + ) # Initialize seed for multi-GPU consistency self._global_seed = 42 @@ -1844,10 +1849,11 @@ 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). + # Reject unsupported top-p-decay and penalty combinations at admission, so + # only the offending request fails (raising later, inside setup_sampler_step + # or sampling, would abort the whole executor step). self._top_p_decay.validate_request(request) + self._penalty_handler.validate_request(request) if self._use_beam_search: if request.py_return_log_probs: if request.py_num_logprobs > 1: @@ -1896,6 +1902,7 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: self._prev_first_finish_reasons_host[slot] = None self._request_grouper.prepare_for_new_request(request, slot) + self._penalty_handler.prepare_for_new_request(request, slot) max_lens = self._finish_reasons_handler.new_max_lens end_ids = self._finish_reasons_handler.new_end_ids @@ -1926,6 +1933,10 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: new_requests, new_seq_slots_cuda_long=seq_slots_tensor_cuda_long ) + self._penalty_handler.update_for_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 @@ -2558,6 +2569,7 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None: else: return None + finalized_token_updates: list[tuple[int, list[int]]] = [] # Fast-path (batched pybind): when the batch is greedy with no beam # search, no logprobs, no draft tokens, no stop-words, and no # speculative tree, collapse per-request pybind chatter into one @@ -2648,6 +2660,7 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None: req.py_rewind_len = 0 else: processed = 1 + num_tokens_before = req.get_num_tokens(DEFAULT_BEAM_IDX) num_accepted = self.process_draft_tokens( req, new_tokens_tensor=new_tokens, @@ -2662,6 +2675,12 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None: req.py_num_accepted_draft_tokens = 0 req.py_rewind_len = 0 processed += num_accepted + if actual_draft_len > 0: + num_new_tokens = req.get_num_tokens(DEFAULT_BEAM_IDX) - num_tokens_before + if num_new_tokens > 0: + assert req.py_seq_slot is not None + confirmed_tokens = req.get_tokens(DEFAULT_BEAM_IDX)[-num_new_tokens:] + finalized_token_updates.append((req.py_seq_slot, confirmed_tokens)) self.handle_logprobs(req, logprobs_state_list=logprobs_state_list, count=processed) req.py_decoding_iter += 1 # Check None or empty list @@ -2672,6 +2691,8 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None: if req.state == LlmRequestState.GENERATION_COMPLETE: self._top_p_decay.retire_slot(req) + self._penalty_handler.update_token_counts(finalized_token_updates) + def _return_log_probs(self, requests: list[LlmRequest]) -> bool: return any(req.py_return_log_probs for req in requests) @@ -3608,6 +3629,20 @@ def _process_requests( logits_cuda, sampling_requests, sampling_requests_metadata.req_num_steps ) + # Apply repetition/presence/frequency penalties in place, before the greedy fast + # path, so both greedy and grouped-sampling logits are penalized. + self._penalty_handler.apply( + logits_cuda, + sampling_requests, + new_tokens=new_tokens_cuda, + seq_slots=seq_slots_cuda, + request_offsets=sampling_requests_metadata.req_offsets, + request_num_steps=sampling_requests_metadata.req_num_steps, + # _is_draft_batch reads requests[0]; an empty batch has no penalties to apply + # anyway, so short-circuit rather than index into it. + is_draft_batch=bool(sampling_requests) and self._is_draft_batch(sampling_requests), + ) + has_min_length = any(getattr(r, "py_min_length", None) for r in sampling_requests) has_bad_words = any(getattr(r, "py_bad_words", None) for r in sampling_requests) # Normalized in executor_request_to_llm_request: a positive int, or diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index aee432e33528..e4bcd807415d 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -15,6 +15,7 @@ l0_a10: tests: # ------------- PyTorch tests --------------- - unittest/_torch/sampler/test_torch_sampler.py + - unittest/_torch/sampler/test_penalties.py - unittest/_torch/test_torch_multi_arange.py - unittest/utils/test_util.py - unittest/utils/test_logger.py diff --git a/tests/unittest/_torch/sampler/test_penalties.py b/tests/unittest/_torch/sampler/test_penalties.py new file mode 100644 index 000000000000..6b9826425ddd --- /dev/null +++ b/tests/unittest/_torch/sampler/test_penalties.py @@ -0,0 +1,515 @@ +# 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. + +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import Fusions +from tensorrt_llm._torch.pyexecutor.sampler.penalties import PenaltyHandler + +apply_batched_occurrence_penalties = Fusions.apply_batched_occurrence_penalties +update_occurrence_workspace = Fusions.update_occurrence_workspace + + +@pytest.fixture(autouse=True) +def _dynamo_recompile_headroom(): + """Recompile headroom for the fullgraph=True penalty op. + + These cases sweep tensor shapes/dtypes, so the op legitimately builds one graph per shape + -- more than the default recompile_limit (8). A served model has fixed shapes; raising the + limit only here avoids tripping fullgraph's hard-fail without touching production. + """ + import torch._dynamo + + with torch._dynamo.config.patch(recompile_limit=128): + yield + + +def _col(values: list[float]) -> torch.Tensor: + return torch.tensor(values, dtype=torch.float32, device="cuda").view(-1, 1) + + +def _dense_penalty_reference( + logits: torch.Tensor, + counts: torch.Tensor, + presence: torch.Tensor | None, + rep: torch.Tensor, + pre: torch.Tensor, + freq: torch.Tensor, + temp: torch.Tensor, +) -> torch.Tensor: + """Dense post-temperature reference for ``apply_batched_occurrence_penalties``. + + Follows the TorchSampler order: repetition where the token is present anywhere + (``counts > 0`` or the prefix mask), then presence + frequency where counted + (``counts > 0``), followed by temperature division in the sampling strategy. + ``rep/pre/freq/temp`` are per-row ``[A, 1]`` tensors. + """ + penalized = logits.float() + present = counts > 0 + if presence is not None: + present = present | (presence > 0) + penalized = torch.where( + present, + torch.where(penalized < 0, penalized * rep, penalized / rep), + penalized, + ) + counts_f = counts.to(torch.float32) + sub = torch.where(counts > 0, pre + freq * counts_f, penalized.new_zeros(())) + return (penalized - sub) / temp + + +def _dense_presence_prefix(counts: torch.Tensor, presence: torch.Tensor) -> torch.Tensor: + prefix = torch.zeros( + presence.size(0), + presence.size(1), + dtype=torch.bool, + device=presence.device, + ) + prefix_slots, prefix_tokens = torch.nonzero(presence, as_tuple=True) + empty = torch.empty(0, dtype=torch.int64, device=presence.device) + update_occurrence_workspace( + counts, + prefix, + empty, + empty, + prefix_slots, + prefix_tokens, + ) + return prefix + + +@pytest.mark.parametrize( + "name,rep,pre,freq,temp,use_prefix", + [ + # repetition only, exercises the sign branch (>1, <1) at temp=1 + ("repetition", [1.3, 2.0, 0.7], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [1.0, 1.0, 1.0], False), + # presence only + ("presence", [1.0, 1.0], [0.5, 1.5], [0.0, 0.0], [1.0, 1.0], False), + # frequency only (counts > 1 -> proportional) + ("frequency", [1.0, 1.0], [0.0, 0.0], [0.4, 0.9], [1.0, 1.0], False), + # combined with temperature != 1 (exercises penalty-before-temperature order) + ( + "combined_temp", + [1.2, 0.8, 1.5], + [0.3, 0.0, 0.7], + [0.2, 0.5, 0.0], + [0.7, 1.3, 2.0], + False, + ), + # ignored-prompt-prefix mask affects repetition only, not presence/frequency + ("prefix", [1.4, 1.1, 0.9], [0.4, 0.6, 0.2], [0.3, 0.1, 0.5], [1.0, 0.8, 1.6], True), + ], +) +@pytest.mark.parametrize("num_steps", [1, 3], ids=["regular", "speculative"]) +def test_penalties_match_dense_logits_reference( + name: str, + rep: list[float], + pre: list[float], + freq: list[float], + temp: list[float], + use_prefix: bool, + num_steps: int, +) -> None: + # vocab=5000 is deliberately not a round power of two. + A, V = len(rep), 5000 + gen = torch.Generator(device="cuda").manual_seed(sum(name.encode()) + num_steps) + logits = torch.randn(A * num_steps, V, device="cuda", generator=gen) * 5.0 + counts = torch.randint(0, 4, (A, V), dtype=torch.int32, device="cuda", generator=gen) + presence = ( + torch.randint(0, 2, (A, V), dtype=torch.int32, device="cuda", generator=gen) + if use_prefix + else None + ) + presence_prefix = _dense_presence_prefix(counts, presence) if presence is not None else None + rep_t, pre_t, freq_t, temp_t = _col(rep), _col(pre), _col(freq), _col(temp) + slots = torch.arange(A, dtype=torch.int64, device="cuda") + row_slots = slots.repeat_interleave(num_steps) + + got = logits.clone() + apply_batched_occurrence_penalties( + got, + counts, + presence_prefix, + torch.ones(A, dtype=torch.bool, device="cuda"), + torch.zeros(A, dtype=torch.bool, device="cuda"), + torch.zeros(1, A, 1, dtype=torch.int32, device="cuda"), + slots, + torch.arange(0, A * num_steps, num_steps, dtype=torch.int32, device="cuda"), + torch.full((A,), num_steps, dtype=torch.int32, device="cuda"), + rep_t.squeeze(1), + pre_t.squeeze(1), + freq_t.squeeze(1), + ) + row_presence = presence[row_slots] if presence is not None else None + ref = _dense_penalty_reference( + logits, + counts[row_slots], + row_presence, + rep_t[row_slots], + pre_t[row_slots], + freq_t[row_slots], + temp_t[row_slots], + ) + # the kernel is pre-temperature-division; divide by temp to compare to the final value. + torch.testing.assert_close(got / temp_t[row_slots], ref, rtol=1e-4, atol=1e-4) + + +def test_penalties_indirect_indexing_bf16() -> None: + # Permuted request offsets and sequence slots penalize a subset of logits rows, with + # repeated slot mappings. Other rows must stay untouched. bfloat16 also covers the + # fp32-compute -> bf16-store cast path. + gen = torch.Generator(device="cuda").manual_seed(3) + num_slots, num_rows, vocab = 5, 10, 3000 + logits = (torch.randn(num_rows, vocab, device="cuda", generator=gen) * 3).to(torch.bfloat16) + orig = logits.clone() + counts = torch.randint( + 0, 4, (num_slots, vocab), dtype=torch.int32, device="cuda", generator=gen + ) + rep = torch.empty(num_slots, device="cuda").uniform_(0.7, 1.6, generator=gen) + pre = torch.empty(num_slots, device="cuda").uniform_(0.0, 0.6, generator=gen) + freq = torch.empty(num_slots, device="cuda").uniform_(0.0, 0.4, generator=gen) + temp = torch.empty(num_slots, device="cuda").uniform_(0.6, 1.4, generator=gen) + # Explicitly exercise permuted rows and repeated slot mappings. + active_rows = torch.tensor([8, 1, 6, 3, 9, 0, 5], dtype=torch.int64, device="cuda") + row_slots = torch.tensor([4, 1, 4, 0, 2, 1, 3], dtype=torch.int64, device="cuda") + + active = torch.ones(num_slots, dtype=torch.bool, device="cuda") + active[1] = False + apply_batched_occurrence_penalties( + logits, + counts, + None, + active, + torch.zeros(num_slots, dtype=torch.bool, device="cuda"), + torch.zeros(1, num_slots, 1, dtype=torch.int32, device="cuda"), + row_slots, + active_rows.to(torch.int32), + torch.ones(active_rows.numel(), dtype=torch.int32, device="cuda"), + rep, + pre, + freq, + ) + + active_row_mask = active[row_slots] + active_slots = row_slots[active_row_mask] + ref = _dense_penalty_reference( + orig[active_rows[active_row_mask]], + counts[active_slots], + None, + rep[active_slots].view(-1, 1), + pre[active_slots].view(-1, 1), + freq[active_slots].view(-1, 1), + temp[active_slots].view(-1, 1), + ) + expected = orig[active_rows].clone() + active_temperature = temp[active_slots].view(-1, 1) + # Recover the pre-temperature op output, then match its fp32-compute -> bf16-store + # boundary. This keeps the tolerance about the op's fp32 math, not bf16 rounding. + expected[active_row_mask] = (ref * active_temperature).to(torch.bfloat16) + torch.testing.assert_close(logits[active_rows], expected, rtol=5e-3, atol=5e-3) + torch.testing.assert_close( + logits[active_rows[~active_row_mask]], + orig[active_rows[~active_row_mask]], + rtol=0, + atol=0, + ) + untouched = torch.ones(num_rows, dtype=torch.bool, device="cuda") + untouched[active_rows] = False + torch.testing.assert_close(logits[untouched], orig[untouched], rtol=0, atol=0) + + +def test_prefix_marking_matches_dense_logits_reference() -> None: + vocab = 70 + counts = torch.zeros(1, vocab, dtype=torch.int32, device="cuda") + presence_prefix = torch.zeros(1, vocab, dtype=torch.bool, device="cuda") + + counted_tokens = torch.tensor([31, 31, 45], dtype=torch.int64, device="cuda") + prefix_tokens = torch.tensor([0, 31, 31, 32, 63, 69], dtype=torch.int64, device="cuda") + counted_slots = torch.zeros_like(counted_tokens) + prefix_slots = torch.zeros_like(prefix_tokens) + update_occurrence_workspace( + counts, + presence_prefix, + counted_slots, + counted_tokens, + prefix_slots, + prefix_tokens, + ) + + logits = torch.linspace(-7.0, 7.0, vocab, device="cuda").view(1, -1) + original = logits.clone() + apply_batched_occurrence_penalties( + logits, + counts, + presence_prefix, + torch.ones(1, dtype=torch.bool, device="cuda"), + torch.zeros(1, dtype=torch.bool, device="cuda"), + torch.zeros(1, 1, 1, dtype=torch.int32, device="cuda"), + torch.zeros(1, dtype=torch.int64, device="cuda"), + torch.zeros(1, dtype=torch.int32, device="cuda"), + torch.ones(1, dtype=torch.int32, device="cuda"), + torch.tensor([1.2], device="cuda"), + torch.tensor([0.4], device="cuda"), + torch.tensor([0.3], device="cuda"), + ) + + dense_prefix = torch.zeros_like(counts) + dense_prefix[0, torch.unique(prefix_tokens)] = 1 + expected = _dense_penalty_reference( + original, + counts, + dense_prefix, + torch.tensor([[1.2]], device="cuda"), + torch.tensor([[0.4]], device="cuda"), + torch.tensor([[0.3]], device="cuda"), + torch.ones(1, 1, device="cuda"), + ) + assert presence_prefix.shape == (1, vocab) + torch.testing.assert_close(logits, expected, rtol=1e-4, atol=1e-4) + + +def test_penalty_op_does_not_latch_pending_token() -> None: + """The penalty op must not write ``has_previous_token``. + + The op reads the flag to decide whether to fold the pending ``new_tokens`` token + into ``counts``; it must never write it (the host re-arms the flag after the op). + Here the flag is False with a stale token far up the vocab: nothing may be folded, + the flag must stay False, and the logits must be untouched. + """ + vocab = 3000 + stale_token = 2500 # a stale pending token far up the vocab + has_previous_token = torch.zeros(1, dtype=torch.bool, device="cuda") + new_tokens = torch.zeros(1, 1, 1, dtype=torch.int32, device="cuda") + new_tokens[0, 0, 0] = stale_token + counts = torch.zeros(1, vocab, dtype=torch.int32, device="cuda") + logits = torch.linspace(-4.0, 4.0, steps=vocab, device="cuda").view(1, vocab) + original = logits.clone() + + apply_batched_occurrence_penalties( + logits, + counts, + None, + torch.ones(1, dtype=torch.bool, device="cuda"), + has_previous_token, + new_tokens, + torch.zeros(1, dtype=torch.int64, device="cuda"), + torch.zeros(1, dtype=torch.int32, device="cuda"), + torch.ones(1, dtype=torch.int32, device="cuda"), + torch.tensor([1.5], device="cuda"), + torch.tensor([0.5], device="cuda"), + torch.tensor([0.4], device="cuda"), + ) + + # Deterministic: the penalty op must leave the latch untouched (host re-arms it). + assert not bool(has_previous_token.item()) + # With has_previous_token False and counts all zero, no penalty may be applied; the + # stale token in particular must not be folded (would perturb logits[2500]). + torch.testing.assert_close(logits, original, rtol=0, atol=0) + + +def _make_handler_request( + *, + slot: int, + tokens: list[int], + prompt_ignore_length: int = 0, + beam_width: int = 1, +) -> SimpleNamespace: + return SimpleNamespace( + sampling_config=SimpleNamespace( + repetition_penalty=[1.2], + presence_penalty=[0.4], + frequency_penalty=[0.3], + temperature=[1.0], + prompt_ignore_length=[prompt_ignore_length], + beam_width=beam_width, + beam_width_array=None, + ), + py_orig_prompt_len=len(tokens), + py_seq_slot=slot, + py_return_log_probs=False, + get_tokens=lambda _beam_idx: tokens, + py_is_draft=False, + ) + + +def _admit(handler: PenaltyHandler, request: SimpleNamespace, slot: int) -> None: + """Admit one request, mirroring TorchSampler.setup_sampler_step. + + ``prepare_for_new_request`` only accumulates on the host; the device buffers are + written by the batched ``update_for_new_requests`` flush at the end of the step. + """ + handler.prepare_for_new_request(request, slot=slot) + handler.update_for_new_requests( + new_seq_slots_cuda_long=torch.tensor([slot], dtype=torch.int64, device="cuda") + ) + + +def _apply_handler( + handler: PenaltyHandler, + request: SimpleNamespace, + logits: torch.Tensor, + num_steps: int, + new_tokens: torch.Tensor, +) -> None: + handler.apply( + logits, + [request], + new_tokens=new_tokens, + seq_slots=torch.tensor([request.py_seq_slot], dtype=torch.int64, device="cuda"), + request_offsets=torch.zeros(1, dtype=torch.int32), + request_num_steps=torch.tensor([num_steps], dtype=torch.int32), + ) + + +def test_handler_tracks_overlap_and_commits_speculative_tail() -> None: + vocab = 16 + slot = 2 + handler = PenaltyHandler( + max_num_sequences=3, + device="cuda", + ) + history = [3] + request = _make_handler_request(slot=slot, tokens=history) + _admit(handler, request, slot) + new_tokens = torch.zeros(3, 3, 1, dtype=torch.int32, device="cuda") + + # The first apply initializes the prompt and marks the first sampled token as + # pending. The request's host history need not be updated before the next apply. + _apply_handler(handler, request, torch.zeros(1, vocab, device="cuda"), 1, new_tokens) + new_tokens[0, slot, 0] = 5 + overlap_logits = torch.linspace(-2.0, 2.0, vocab, device="cuda").view(1, vocab) + overlap_original = overlap_logits.clone() + _apply_handler(handler, request, overlap_logits, 1, new_tokens) + overlap_counts = torch.bincount(torch.tensor([3, 5], device="cuda"), minlength=vocab).to( + torch.int32 + )[None] + overlap_expected = _dense_penalty_reference( + overlap_original, + overlap_counts, + None, + torch.full((1, 1), 1.2, device="cuda"), + torch.full((1, 1), 0.4, device="cuda"), + torch.full((1, 1), 0.3, device="cuda"), + torch.ones(1, 1, device="cuda"), + ) + torch.testing.assert_close(overlap_logits, overlap_expected, rtol=1e-4, atol=1e-4) + + # The next invocation is speculative. All rows use the same confirmed history; + # the current draft window remains tentative until acceptance is resolved. + history.extend([5, 6]) + new_tokens[0, slot, 0] = 6 + spec_logits = torch.linspace(-3.0, 3.0, steps=3 * vocab, device="cuda").view(3, vocab) + spec_original = spec_logits.clone() + _apply_handler(handler, request, spec_logits, 3, new_tokens) + spec_counts = torch.bincount(torch.tensor(history, device="cuda"), minlength=vocab).to( + torch.int32 + )[None] + spec_expected = _dense_penalty_reference( + spec_original, + spec_counts.expand(3, -1), + None, + torch.full((3, 1), 1.2, device="cuda"), + torch.full((3, 1), 0.4, device="cuda"), + torch.full((3, 1), 0.3, device="cuda"), + torch.ones(3, 1, device="cuda"), + ) + torch.testing.assert_close(spec_logits, spec_expected, rtol=1e-4, atol=1e-4) + + # Sampler-side acceptance commits the complete finalized sequence. Deliberately + # leave a different raw target token in the device buffer, as rejection sampling + # can do; clearing the pending flag must prevent it from entering the workspace. + history.extend([7, 8, 7]) + new_tokens[0, slot, 0] = 4 + handler.update_token_counts([(slot, [7, 8, 7])]) + logits = torch.linspace(-4.0, 4.0, steps=3 * vocab, device="cuda").view(3, vocab) + original = logits.clone() + _apply_handler(handler, request, logits, 3, new_tokens) + + expected_counts = torch.bincount(torch.tensor(history, device="cuda"), minlength=vocab).to( + torch.int32 + )[None] + expected = _dense_penalty_reference( + original, + expected_counts.expand(3, -1), + None, + torch.full((3, 1), 1.2, device="cuda"), + torch.full((3, 1), 0.4, device="cuda"), + torch.full((3, 1), 0.3, device="cuda"), + torch.ones(3, 1, device="cuda"), + ) + torch.testing.assert_close(logits, expected, rtol=1e-4, atol=1e-4) + + +def test_regular_handler_slot_reuse_does_not_leak_penalties() -> None: + vocab = 16 + handler = PenaltyHandler( + max_num_sequences=1, + device="cuda", + ) + new_tokens = torch.zeros(1, 1, 1, dtype=torch.int32, device="cuda") + + first = _make_handler_request(slot=0, tokens=[3, 3], prompt_ignore_length=1) + _admit(handler, first, 0) + _apply_handler(handler, first, torch.zeros(1, vocab, device="cuda"), 1, new_tokens) + + second = _make_handler_request(slot=0, tokens=[5]) + _admit(handler, second, 0) + logits = torch.linspace(-2.0, 2.0, steps=vocab, device="cuda").view(1, vocab) + original = logits.clone() + _apply_handler(handler, second, logits, 1, new_tokens) + + expected_counts = torch.zeros(1, vocab, dtype=torch.int32, device="cuda") + expected_counts[0, 5] = 1 + expected = _dense_penalty_reference( + original, + expected_counts, + None, + torch.full((1, 1), 1.2, device="cuda"), + torch.full((1, 1), 0.4, device="cuda"), + torch.full((1, 1), 0.3, device="cuda"), + torch.ones(1, 1, device="cuda"), + ) + torch.testing.assert_close(logits, expected, rtol=1e-4, atol=1e-4) + + +def test_handler_ignores_occurrence_penalties_with_beam_search() -> None: + """Beam-search requests never become penalty-active. + + ``PenaltyHandler.validate_request`` rejects this combination at admission, so the + handler should only ever see beam_width == 1 requests. It stays defensive anyway: + a beam-search request leaves its slot inactive, and ``apply`` is then a no-op. + """ + vocab = 16 + handler = PenaltyHandler(max_num_sequences=1, device="cuda") + new_tokens = torch.zeros(1, 1, 1, dtype=torch.int32, device="cuda") + + request = _make_handler_request(slot=0, tokens=[3, 3], beam_width=2) + _admit(handler, request, 0) + + logits = torch.linspace(-2.0, 2.0, steps=vocab, device="cuda").view(1, vocab) + original = logits.clone() + _apply_handler(handler, request, logits, 1, new_tokens) + + assert not bool(handler.store.active_cuda[0].item()) + torch.testing.assert_close(logits, original, rtol=0, atol=0) + + +def test_validate_request_rejects_penalties_with_beam_search() -> None: + """The admission-time check that keeps the combination above from arriving.""" + PenaltyHandler.validate_request(_make_handler_request(slot=0, tokens=[3], beam_width=1)) + with pytest.raises(ValueError, match="penalties with beam search"): + PenaltyHandler.validate_request(_make_handler_request(slot=0, tokens=[3], beam_width=2)) diff --git a/tests/unittest/_torch/sampler/test_penalties_e2e.py b/tests/unittest/_torch/sampler/test_penalties_e2e.py new file mode 100644 index 000000000000..332f8a675c14 --- /dev/null +++ b/tests/unittest/_torch/sampler/test_penalties_e2e.py @@ -0,0 +1,329 @@ +# 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. + +from dataclasses import dataclass +from pathlib import Path + +import pytest +import torch +from utils.llm_data import llm_models_root + +from tensorrt_llm import LLM, SamplingParams +from tensorrt_llm.executor.result import CompletionOutput, GenerationResult +from tensorrt_llm.llmapi import CudaGraphConfig, NGramDecodingConfig +from tensorrt_llm.llmapi import KvCacheConfig as TRT_KvCacheConfig + + +@pytest.fixture(scope="module") +def model_path() -> Path: + return llm_models_root() / "llama-models-v2/TinyLlama-1.1B-Chat-v1.0" + + +@dataclass(frozen=True) +class _PenaltyE2ECase: + name: str + prompt: str + sampling_params: SamplingParams + + +def _penalty_sampling_params( + max_tokens: int = 1, + logprobs: int = 1, + **penalties: float | int, +) -> SamplingParams: + return SamplingParams( + max_tokens=max_tokens, + temperature=1.3, + seed=12345, + ignore_eos=True, + logprobs=logprobs, + logprobs_mode="processed", + return_generation_logits=True, + **penalties, + ) + + +def _make_penalty_e2e_cases() -> list[_PenaltyE2ECase]: + repeated_answer_prompt = "The capital of France is Paris. The capital of France is" + capital_prompt = "The capital of France is" + repeated_token_prompt = "cat cat cat cat The capital of France is" + + return [ + _PenaltyE2ECase( + "repetition_discourage", + repeated_answer_prompt, + _penalty_sampling_params(repetition_penalty=100.0), + ), + _PenaltyE2ECase( + "repetition_encourage", + capital_prompt, + _penalty_sampling_params(repetition_penalty=0.01), + ), + _PenaltyE2ECase( + "additive_reward", + repeated_token_prompt, + _penalty_sampling_params(presence_penalty=-10.0, frequency_penalty=-2.0), + ), + _PenaltyE2ECase( + "frequency_count", + repeated_token_prompt, + _penalty_sampling_params(frequency_penalty=5.0), + ), + _PenaltyE2ECase( + "additive_prompt_ignored", + capital_prompt, + _penalty_sampling_params( + presence_penalty=100.0, + frequency_penalty=100.0, + prompt_ignore_length=10_000, + ), + ), + _PenaltyE2ECase( + "combined_penalties", + repeated_answer_prompt, + _penalty_sampling_params( + max_tokens=6, + logprobs=5, + repetition_penalty=1.7, + presence_penalty=2.0, + frequency_penalty=0.75, + prompt_ignore_length=2, + ), + ), + ] + + +def _create_torch_llm( + model_dir: Path, + max_batch_size: int | None = None, + speculative_config: NGramDecodingConfig | None = None, + enable_iter_perf_stats: bool = False, +) -> LLM: + llm_kwargs: dict[str, object] = {} + if max_batch_size is not None: + llm_kwargs["max_batch_size"] = max_batch_size + if speculative_config is not None: + llm_kwargs["speculative_config"] = speculative_config + + return LLM( + model=str(model_dir), + tensor_parallel_size=1, + trust_remote_code=True, + enable_chunked_prefill=True, + cuda_graph_config=CudaGraphConfig(), + sampler_type="TorchSampler", + kv_cache_config=TRT_KvCacheConfig(enable_block_reuse=False), + max_num_tokens=128, + enable_iter_perf_stats=enable_iter_perf_stats, + **llm_kwargs, + ) + + +def _run_penalty_e2e_cases( + model_dir: Path, + cases: list[_PenaltyE2ECase], +) -> tuple[dict[str, GenerationResult], dict[str, tuple[int, ...]]]: + with _create_torch_llm(model_dir) as llm: + outputs = llm.generate( + [case.prompt for case in cases], + sampling_params=[case.sampling_params for case in cases], + use_tqdm=False, + ) + + results = dict(zip((case.name for case in cases), outputs, strict=True)) + prompt_token_ids = { + case.name: tuple(int(token_id) for token_id in output.prompt_token_ids) + for case, output in zip(cases, outputs, strict=True) + } + return results, prompt_token_ids + + +def _reference_penalized_logits( + raw_logits: torch.Tensor, + token_history: list[int], + prompt_length: int, + sampling_params: SamplingParams, +) -> torch.Tensor: + """Apply the documented penalties independently of TorchSampler.""" + vocab_size = raw_logits.numel() + history = torch.tensor(token_history, dtype=torch.int64) + valid_history = history[(history >= 0) & (history < vocab_size)] + adjusted_logits = raw_logits.float() + + repetition_penalty = sampling_params.repetition_penalty or 1.0 + if repetition_penalty != 1.0 and valid_history.numel() > 0: + repetition_mask = torch.bincount(valid_history, minlength=vocab_size).bool() + repetition_scaled_logits = torch.where( + adjusted_logits < 0, + adjusted_logits * repetition_penalty, + adjusted_logits / repetition_penalty, + ) + adjusted_logits = torch.where(repetition_mask, repetition_scaled_logits, adjusted_logits) + + prompt_ignore_length = sampling_params.prompt_ignore_length or 0 + occurrence_start = max(0, min(prompt_ignore_length, prompt_length)) + occurrence_history = history[occurrence_start:] + valid_occurrences = occurrence_history[ + (occurrence_history >= 0) & (occurrence_history < vocab_size) + ] + occurrence_counts = torch.bincount(valid_occurrences, minlength=vocab_size).float() + + presence_penalty = sampling_params.presence_penalty or 0.0 + frequency_penalty = sampling_params.frequency_penalty or 0.0 + adjusted_logits -= presence_penalty * (occurrence_counts > 0) + adjusted_logits -= frequency_penalty * occurrence_counts + + dtype_limit = torch.finfo(raw_logits.dtype).max + return adjusted_logits.clamp(min=-dtype_limit, max=dtype_limit).to(raw_logits.dtype) + + +def _reference_processed_logprobs( + raw_logits: torch.Tensor, + token_history: list[int], + prompt_length: int, + sampling_params: SamplingParams, +) -> torch.Tensor: + penalized_logits = _reference_penalized_logits( + raw_logits, + token_history, + prompt_length, + sampling_params, + ) + temperature = sampling_params.temperature + if temperature is not None and temperature != 0.0: + penalized_logits = penalized_logits / max(temperature, 1e-5) + processed_logits = penalized_logits.float() + sampling_probs = torch.softmax(processed_logits, dim=-1) + processed_logits = processed_logits.masked_fill(sampling_probs == 0, float("-inf")) + return torch.log_softmax(processed_logits, dim=-1) + + +def _assert_completion_penalty_logprobs( + case: _PenaltyE2ECase, + completion: CompletionOutput, + prompt_token_ids: tuple[int, ...], +) -> None: + assert completion.token_ids is not None, case.name + assert completion.generation_logits is not None, case.name + assert completion.logprobs is not None, case.name + + token_history = list(prompt_token_ids) + expected_cumulative_logprob = 0.0 + for step, (token_id, raw_logits, actual_logprobs) in enumerate( + zip(completion.token_ids, completion.generation_logits, completion.logprobs, strict=True) + ): + location = f"{case.name}/step_{step}" + assert token_id in actual_logprobs, location + expected_logprobs = _reference_processed_logprobs( + raw_logits, + token_history, + len(prompt_token_ids), + case.sampling_params, + ) + + for returned_token_id, actual in actual_logprobs.items(): + assert actual.logprob == pytest.approx( + float(expected_logprobs[returned_token_id]), + rel=2e-5, + abs=2e-4, + ), location + + num_logprobs = case.sampling_params.logprobs + if num_logprobs: + ranked_logprobs = { + actual.rank: actual.logprob + for actual in actual_logprobs.values() + if actual.rank is not None and actual.rank <= num_logprobs + } + assert set(ranked_logprobs) == set(range(1, num_logprobs + 1)), location + expected_top_logprobs = torch.topk(expected_logprobs, k=num_logprobs).values + for rank, expected in enumerate(expected_top_logprobs, start=1): + assert ranked_logprobs[rank] == pytest.approx( + float(expected), rel=2e-5, abs=2e-4 + ), location + + expected_cumulative_logprob += float(expected_logprobs[token_id]) + token_history.append(token_id) + + assert completion.cumulative_logprob == pytest.approx( + expected_cumulative_logprob, rel=2e-5, abs=2e-4 + ), case.name + + +@pytest.mark.high_cuda_memory +def test_torch_sampler_penalty_logits_e2e(model_path: Path) -> None: + """Validate TorchSampler's processed logits against the penalty formulas.""" + cases = _make_penalty_e2e_cases() + results, prompt_token_ids = _run_penalty_e2e_cases(model_path, cases) + + for case in cases: + for completion in results[case.name].outputs: + _assert_completion_penalty_logprobs( + case, + completion, + prompt_token_ids[case.name], + ) + + +@pytest.mark.high_cuda_memory +def test_torch_sampler_speculative_penalty_e2e(model_path: Path) -> None: + """Validate the speculative (NGram) path's penalized logprobs against the formula. + + A positive temperature keeps the processed logprobs a real distribution the penalty + formula can be checked against (greedy would collapse them to one-hot). At this + temperature the penalties push the target away from NGram's repetition drafts, so drafts + are proposed but not accepted -- every emitted token is target-sampled and its processed + logprobs match the formula. Accepted draft tokens report a one-hot logprob and cannot be + formula-checked; the accepted-token confirmed-history commit path is covered at the logit + level by ``test_handler_tracks_overlap_and_commits_speculative_tail``. + """ + case = _PenaltyE2ECase( + "ngram_speculative_penalties", + "red blue red blue red blue red blue red blue", + _penalty_sampling_params( + max_tokens=8, + logprobs=5, + repetition_penalty=1.5, + presence_penalty=1.0, + frequency_penalty=1.0, + prompt_ignore_length=2, + ), + ) + speculative_config = NGramDecodingConfig( + max_draft_len=3, + max_matching_ngram_size=2, + is_keep_all=True, + is_use_oldest=True, + is_public_pool=False, + ) + with _create_torch_llm( + model_path, + max_batch_size=1, + speculative_config=speculative_config, + enable_iter_perf_stats=True, + ) as llm: + speculative_outputs = llm.generate( + [case.prompt], sampling_params=[case.sampling_params], use_tqdm=False + ) + stats = llm.get_stats(timeout=5) + + assert any(stat.get("specDecodingStats", {}).get("numDraftTokens", 0) > 0 for stat in stats), ( + "NGram must produce draft tokens in this test" + ) + + speculative_prompt_token_ids = tuple( + int(token_id) for token_id in speculative_outputs[0].prompt_token_ids + ) + for completion in speculative_outputs[0].outputs: + _assert_completion_penalty_logprobs(case, completion, speculative_prompt_token_ids)