diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/trtllm_triton.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/trtllm_triton.py new file mode 100644 index 000000000000..a08c17fcbeb0 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/sampler/ops/trtllm_triton.py @@ -0,0 +1,236 @@ +# 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. + +"""TRT-LLM-internal Triton kernels for the sampler ops layer. + +Triton kernels for the occurrence penalties (repetition / presence / frequency), the +torch/Triton counterpart of the C++ ``batchApplyPenalty`` kernel +(``cpp/tensorrt_llm/kernels/penaltyKernels.cu``), with no dependency on the +sampling_utils interface or other backend modules: + +* :func:`update_occurrence_workspace` -- initializes prompt occurrence state. +* :func:`apply_batched_occurrence_penalties_triton` -- increments regular state and + applies penalties in one pass over packed logits. +""" + +import torch +import triton +import triton.language as tl + + +def update_occurrence_workspace( + counts: torch.Tensor, + presence_prefix: torch.Tensor | None, + counted_slots: torch.Tensor, + counted_tokens: torch.Tensor, + prefix_slots: torch.Tensor, + prefix_tokens: torch.Tensor, +) -> None: + """Fold newly-committed tokens into the persistent occurrence workspace. + + Mirrors the logical accumulation ``batchApplyPenalty`` performs in its + ``penaltyWorkspace`` while packing the prefix-presence half: tokens in the ignored + prompt prefix ``[0, prompt_ignore_length)`` set the presence bitmap only (they count + for repetition but not for presence/frequency), while all other tokens (the rest of + the prompt plus every generated token) increment the occurrence counts. + + Arguments (all index tensors are 1-D and pre-split on the host): + counts: ``int32[num_slots, vocab_size]`` occurrence counts, updated in place. + presence_prefix: packed int32 ``[num_slots, ceil(vocab_size / 32)]`` prefix + presence bitmap, or ``None`` when no active request uses + ``prompt_ignore_length``. + counted_slots / counted_tokens: (slot, token) pairs to increment in ``counts``. + prefix_slots / prefix_tokens: (slot, token) pairs to mark in ``presence_prefix``. + """ + if counted_slots.numel() > 0: + ones = torch.ones(counted_slots.shape[0], dtype=counts.dtype, device=counts.device) + # accumulate=True sums repeated (slot, token) pairs -> occurrence count. + counts.index_put_((counted_slots, counted_tokens), ones, accumulate=True) + if presence_prefix is not None and prefix_slots.numel() > 0: + block_size = 256 + _mark_presence_prefix_kernel[(triton.cdiv(prefix_slots.numel(), block_size),)]( + presence_prefix, + prefix_slots, + prefix_tokens, + presence_prefix.stride(0), + prefix_slots.numel(), + BLOCK_SIZE=block_size, + ) + + +@triton.jit +def _mark_presence_prefix_kernel( + presence_prefix_ptr, + prefix_slots_ptr, + prefix_tokens_ptr, + presence_prefix_row_stride, + num_tokens, + BLOCK_SIZE: tl.constexpr, +) -> None: + """Atomically mark ignored-prefix tokens in the packed presence bitmap.""" + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < num_tokens + slots = tl.load(prefix_slots_ptr + offsets, mask=mask, other=0) + tokens = tl.load(prefix_tokens_ptr + offsets, mask=mask, other=0).to(tl.int32) + word_offsets = tokens // 32 + bit_offsets = tokens % 32 + bits = tl.full((BLOCK_SIZE,), 1, tl.int32) << bit_offsets + tl.atomic_or( + presence_prefix_ptr + slots * presence_prefix_row_stride + word_offsets, + bits, + mask=mask, + ) + + +@triton.jit +def _apply_batched_occurrence_penalties_kernel( + logits_ptr, + counts_ptr, + presence_prefix_ptr, + active_ptr, + has_previous_token_ptr, + new_tokens_ptr, + seq_slots_ptr, + request_offsets_ptr, + request_num_steps_ptr, + repetition_ptr, + presence_ptr, + frequency_ptr, + vocab, + logits_row_stride, + workspace_row_stride, + presence_prefix_row_stride, + new_tokens_step_stride, + new_tokens_slot_stride, + new_tokens_beam_stride, + LOGIT_LIMIT: tl.constexpr, + HAS_PRESENCE_PREFIX: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +) -> None: + """Update the workspace once and apply it to every row of a request.""" + request_idx = tl.program_id(0) + vocab_block = tl.program_id(1) + + slot = tl.load(seq_slots_ptr + request_idx) + active = tl.load(active_ptr + slot) != 0 + num_steps = tl.load(request_num_steps_ptr + request_idx) + if not active or num_steps <= 0: + return + + request_offset = tl.load(request_offsets_ptr + request_idx) + + offsets = vocab_block * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + vocab_mask = offsets < vocab + + count_offset = slot * workspace_row_stride + offsets + count = tl.load(counts_ptr + count_offset, mask=vocab_mask, other=0) + has_previous_token = tl.load(has_previous_token_ptr + slot) != 0 + previous_token = tl.load( + new_tokens_ptr + + slot * new_tokens_slot_stride + + 0 * new_tokens_step_stride + + 0 * new_tokens_beam_stride, + mask=has_previous_token, + other=-1, + ) + count += tl.where(has_previous_token & (offsets == previous_token), 1, 0) + + seen = count > 0 + if HAS_PRESENCE_PREFIX: + prefix_word_offsets = vocab_block * BLOCK_SIZE // 32 + tl.arange(0, BLOCK_SIZE // 32) + prefix_words = tl.load( + presence_prefix_ptr + slot * presence_prefix_row_stride + prefix_word_offsets, + mask=prefix_word_offsets < tl.cdiv(vocab, 32), + other=0, + ) + prefix_seen = (prefix_words[:, None] >> tl.arange(0, 32)[None, :]) & 1 + seen |= prefix_seen.to(tl.int1).reshape(BLOCK_SIZE) + + repetition = tl.load(repetition_ptr + slot) + presence = tl.load(presence_ptr + slot) + frequency = tl.load(frequency_ptr + slot) + + # A request owns one occurrence row. Looping over its packed logits rows in the + # same program ensures every speculative row observes the same updated history; + # separate step programs would race with the count store below. Use the runtime + # request length so step counts do not create separate compiled kernel variants. + for step_idx in tl.range(0, num_steps): + logit_offset = (request_offset + step_idx) * logits_row_stride + offsets + logit = tl.load(logits_ptr + logit_offset, mask=vocab_mask, other=0.0).to(tl.float32) + + repeated = tl.where(logit < 0.0, logit * repetition, logit / repetition) + logit = tl.where(seen, repeated, logit) + logit -= tl.where( + count > 0, + presence + frequency * count.to(tl.float32), + 0.0, + ) + logit = tl.maximum(-LOGIT_LIMIT, tl.minimum(logit, LOGIT_LIMIT)) + tl.store( + logits_ptr + logit_offset, + logit.to(logits_ptr.dtype.element_ty), + mask=vocab_mask, + ) + + tl.store(counts_ptr + count_offset, count, mask=vocab_mask) + + +def apply_batched_occurrence_penalties_triton( + logits: torch.Tensor, + counts: torch.Tensor, + presence_prefix: torch.Tensor | None, + active: torch.Tensor, + has_previous_token: torch.Tensor, + new_tokens: torch.Tensor, + seq_slots: torch.Tensor, + request_offsets: torch.Tensor, + request_num_steps: torch.Tensor, + repetition: torch.Tensor, + presence: torch.Tensor, + frequency: torch.Tensor, +) -> None: + """Apply occurrence penalties to ``logits`` in place, before temperature handling.""" + num_requests = seq_slots.numel() + if num_requests == 0: + return + + vocab = logits.size(-1) + block_size = 1024 + grid = (num_requests, triton.cdiv(vocab, block_size)) + has_presence_prefix = presence_prefix is not None + _apply_batched_occurrence_penalties_kernel[grid]( + logits, + counts, + presence_prefix if has_presence_prefix else counts, + active, + has_previous_token, + new_tokens, + seq_slots, + request_offsets, + request_num_steps, + repetition, + presence, + frequency, + vocab, + logits.stride(0), + counts.stride(0), + presence_prefix.stride(0) if presence_prefix is not None else counts.stride(0), + new_tokens.stride(0), + new_tokens.stride(1), + new_tokens.stride(2), + LOGIT_LIMIT=torch.finfo(logits.dtype).max, + HAS_PRESENCE_PREFIX=has_presence_prefix, + BLOCK_SIZE=block_size, + num_warps=4, + ) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index e7fe112ace6f..64188e77243e 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -88,6 +88,10 @@ from ..llm_request import LlmRequest, LlmRequestState, get_draft_token_length from ..resource_manager import ResourceManager, ResourceManagerType from ..scheduler import ScheduledRequests +from .ops.trtllm_triton import ( + apply_batched_occurrence_penalties_triton, + update_occurrence_workspace, +) from .sampling_utils import ( BEAM_SEARCH_PAD_TOKEN, GREEDY, @@ -448,6 +452,18 @@ def _unwrap_singleton(p: Optional[List[T]]) -> Optional[T]: return t +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) + ) + + def _get_beam_width_in(request: LlmRequest) -> int: return ( 1 @@ -1332,6 +1348,338 @@ def _record_sampler_event( ) +class _OccurrencePenaltyHandler: + """Applies repetition / presence / frequency penalties for :class:`TorchSampler`. + + Occurrence tracking and the individual penalty transforms follow the C++ + ``batchApplyPenalty`` kernel (``cpp/tensorrt_llm/kernels/penaltyKernels.cu``) as + driven by ``PenaltyLayer``. Their ordering follows TorchSampler: penalties are + applied before the sampling strategy handles temperature. All persistent device + state lives in a single :class:`_PenaltyStore`, the torch counterpart of the tensors + ``PenaltyLayer`` allocates: + + * per-slot penalty-parameter buffers (``repetition`` / ``presence`` / ``frequency``, + shape ``[max_num_sequences]``) -- the counterpart of + ``allocateBuffer`` + ``fillBuffers``; filled once per request in + ``prepare_for_new_request`` and gathered per step (no per-step host rebuild); + * an occurrence workspace (dense ``counts`` plus a packed ``presence_prefix`` + bitmap) -- the counterpart of ``allocateWorkspace`` / + ``mPenaltyWorkspaceDevice``; updated incrementally each step. + + The actual logits math lives in the ops module. Parameter buffers are allocated + eagerly (tiny, vocab-independent); vocab-sized workspaces are allocated lazily and + skipped entirely when no matching request uses an occurrence penalty. + """ + + @dataclass(kw_only=True) + class _PenaltyStore: + """Persistent device state: penalty-parameter buffers + occurrence workspace.""" + + # --- Penalty parameters (allocateBuffer counterpart), shape [max_num_sequences] --- + repetition: torch.Tensor + """float32; per-slot repetition penalty (default 1.0).""" + presence: torch.Tensor + """float32; per-slot presence penalty (default 0.0).""" + frequency: torch.Tensor + """float32; per-slot frequency penalty (default 0.0).""" + active: torch.Tensor + """bool[slots]; whether a slot has an active occurrence penalty.""" + has_previous_token: torch.Tensor + """bool[slots]; whether ``new_tokens`` contains a token to accumulate.""" + + # --- Occurrence workspace (allocateWorkspace counterpart), allocated lazily --- + counts: torch.Tensor | None = None + """int32 or None; per-slot occurrence counts of prompt (beyond + prompt_ignore_length) and generated tokens. Drives presence/frequency and (via + counts > 0) repetition.""" + presence_prefix: torch.Tensor | None = None + """Packed int32 or None; per-slot presence bitmap with shape + ``[slots, ceil(vocab / 32)]`` for the ignored prompt prefix + ``[0, prompt_ignore_length)``. It contributes to repetition only.""" + + @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 bitmap. + self._needs_prefix = False + self._num_active_slots = 0 + # Per-slot state; None marks a slot without active occurrence penalties. + self._slots: list[_OccurrencePenaltyHandler._SlotState | None] = [None] * max_num_sequences + # Eagerly allocate the (tiny, vocab-independent) parameter buffers with their + # no-op defaults. inference_mode(False) so they stay mutable in place later. + with torch.inference_mode(False): + self.store = self._PenaltyStore( + repetition=torch.ones(max_num_sequences, dtype=torch.float32, device=self._device), + presence=torch.zeros(max_num_sequences, dtype=torch.float32, device=self._device), + frequency=torch.zeros(max_num_sequences, dtype=torch.float32, device=self._device), + active=torch.zeros(max_num_sequences, dtype=torch.bool, device=self._device), + has_previous_token=torch.zeros( + max_num_sequences, dtype=torch.bool, device=self._device + ), + ) + # Persistent destinations for the per-step metadata consumed by Triton. + self._request_offsets = torch.empty( + max_num_sequences, dtype=torch.int32, device=self._device + ) + self._request_num_steps = torch.empty( + max_num_sequences, dtype=torch.int32, device=self._device + ) + + @staticmethod + def _fill_slot(tensor: torch.Tensor, slot: int, value: bool | float | int) -> None: + """Asynchronously fill one slot without CUDA scalar-assignment synchronization.""" + tensor.narrow(0, slot, 1).fill_(value) + + def prepare_for_new_request(self, request: LlmRequest, slot: int) -> None: + """Fill the slot's penalty parameters and reset its workspace. + + Called from ``TorchSampler.setup_sampler_step`` for each new request, mirroring + ``PenaltyLayer::setup`` (``fillBuffers`` + per-``batchSlot`` ``setZero``). + 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 + self._fill_slot(self.store.active, slot, False) + 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 + + state = self._SlotState(prompt_ignore_length=prompt_ignore_length) + self._slots[slot] = state + # Fill the per-slot parameter buffers (once per request; gathered every step). + self._fill_slot( + self.store.repetition, + slot, + repetition if repetition is not None else 1.0, + ) + self._fill_slot( + self.store.presence, + slot, + presence if presence is not None else 0.0, + ) + self._fill_slot( + self.store.frequency, + slot, + frequency if frequency is not None else 0.0, + ) + if not was_active: + self._num_active_slots += 1 + self._fill_slot(self.store.active, slot, True) + self._fill_slot(self.store.has_previous_token, slot, False) + + # Re-zero the workspace slot so a prior occupant's counts do not leak in. + if self.store.counts is not None: + self.store.counts[slot].zero_() + + if self.store.presence_prefix is not None: + self.store.presence_prefix[slot].zero_() + + def _ensure_workspace(self, vocab_size: int) -> None: + # Lazily allocate the occurrence workspace on first use (vocab_size is only known + # here), mirroring PenaltyLayer::allocateWorkspace being gated on penalty usage. + # inference_mode(False) so the persistent tensors can be mutated in place later. + with torch.inference_mode(False): + if self.store.counts is None: + self.store.counts = torch.zeros( + (self._max_num_sequences, vocab_size), dtype=torch.int32, device=self._device + ) + if self._needs_prefix and self.store.presence_prefix is None: + packed_vocab_size = (vocab_size + 31) // 32 + self.store.presence_prefix = torch.zeros( + (self._max_num_sequences, packed_vocab_size), + dtype=torch.int32, + device=self._device, + ) + + 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 _initialize_workspace( + self, + request: LlmRequest, + state: "_OccurrencePenaltyHandler._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 = self.store.counts + assert counts is not None + counted_tokens: list[int] = [] + prefix_tokens: list[int] = [] + for index, token in enumerate(request.get_tokens(0)[: request.py_orig_prompt_len]): + if token < 0 or token >= vocab_size: + continue + if index < state.prompt_ignore_length: + prefix_tokens.append(token) + else: + counted_tokens.append(token) + + slot_index = torch.full( + (len(counted_tokens),), slot, dtype=torch.int64, device=self._device + ) + prefix_slot_index = torch.full( + (len(prefix_tokens),), slot, dtype=torch.int64, device=self._device + ) + update_occurrence_workspace( + counts, + self.store.presence_prefix, + slot_index, + self._to_device(counted_tokens, torch.int64), + prefix_slot_index, + self._to_device(prefix_tokens, torch.int64), + ) + state.initialized = True + + 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 = self.store.counts + assert counts is not None + vocab_size = counts.size(-1) + counted_slots: list[int] = [] + counted_tokens: list[int] = [] + + for slot, tokens in updates: + if self._slots[slot] is None: + continue + self._fill_slot(self.store.has_previous_token, slot, False) + for token in tokens: + if 0 <= token < vocab_size: + counted_slots.append(slot) + counted_tokens.append(token) + + if not counted_tokens: + return + + update_occurrence_workspace( + counts, + self.store.presence_prefix, + self._to_device(counted_slots, torch.int64), + self._to_device(counted_tokens, torch.int64), + self._to_device([], torch.int64), + self._to_device([], torch.int64), + ) + + @nvtx_range("apply_occurrence_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, + ) -> None: + """Apply occurrence penalties to ``logits`` in place. + + ``logits`` is the packed generated-token logits ``[sum(num_steps * num_beams), + vocab_size]``; the row layout matches ``TorchSampler._apply_min_length_penalty``. + """ + if 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, "_OccurrencePenaltyHandler._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 + + self._ensure_workspace(logits.size(-1)) + store = self.store + counts = store.counts + assert counts is not None + for request, state in active_requests: + self._initialize_workspace(request, state, logits.size(-1)) + + num_requests = len(requests) + kernel_request_offsets = self._request_offsets[:num_requests] + kernel_request_num_steps = self._request_num_steps[:num_requests] + kernel_request_offsets.copy_(request_offsets, non_blocking=True) + kernel_request_num_steps.copy_(request_num_steps, non_blocking=True) + apply_batched_occurrence_penalties_triton( + logits, + counts, + store.presence_prefix, + store.active, + store.has_previous_token, + new_tokens, + seq_slots, + kernel_request_offsets, + kernel_request_num_steps, + store.repetition, + store.presence, + store.frequency, + ) + # Flag has_previous_token for the slots this call penalized (active, num_steps > 0) + # so the next apply folds their sampled new_tokens. Kept on the host, out of the + # multi-vocab-block kernel, to avoid a cross-block race on the flag. + 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.index_fill_( + 0, self._to_device(pending_token_slots, torch.int64), True + ) + + class TorchSampler(Sampler[SampleStateTorch], AsyncWorkerMixin): DEFAULT_MAX_STOP_WORD_LENGTH = 20 DEFAULT_MAX_STOP_WORDS = 10 @@ -2375,6 +2723,10 @@ def __init__(self, args: Args): self.store.new_tokens.shape == self._finish_reasons_handler.store.finish_reasons_cuda.shape ) + self._penalty_handler = _OccurrencePenaltyHandler( + max_num_sequences=self.max_num_sequences, + device="cuda", + ) # Initialize seed for multi-GPU consistency self._global_seed = 42 @@ -2924,6 +3276,11 @@ def _collect_new_requests_for_setup( @override def validate_request(self, request: LlmRequest) -> None: if self._use_beam_search: + if _get_max_beam_width(request) > 1 and _has_occurrence_penalty(request): + logger.warning( + "TorchSampler does not support repetition, presence, or frequency " + "penalties with beam search; these penalties will be ignored." + ) if request.py_return_log_probs: if request.py_num_logprobs > 1: raise ValueError( @@ -2971,6 +3328,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 @@ -3674,6 +4032,7 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None: else: return None + finalized_token_updates: list[tuple[int, list[int]]] = [] for req_idx, req in enumerate(state.requests): if req.state == LlmRequestState.GENERATION_COMPLETE: continue @@ -3721,6 +4080,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, @@ -3735,6 +4095,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 @@ -3743,6 +4109,8 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None: req.py_seq_slot ] = req.py_num_accepted_draft_tokens + 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) @@ -4697,6 +5065,17 @@ def _process_requests( sampling_requests_metadata.req_num_beams.tolist(), ) + # 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, + ) + # Fast path for greedy sampling if self._can_use_fast_greedy_path(sampling_requests): # Compute destination indices on CPU (same pattern as _unbatch_sampling_results) diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 678adad9a4ce..0b985e7f0439 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..3b792f87a986 --- /dev/null +++ b/tests/unittest/_torch/sampler/test_penalties.py @@ -0,0 +1,460 @@ +# 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.trtllm_triton import ( + apply_batched_occurrence_penalties_triton, + update_occurrence_workspace, +) +from tensorrt_llm._torch.pyexecutor.sampler.sampler import _OccurrencePenaltyHandler + + +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_occurrence_penalties_triton``. + + Follows the TorchSampler order: repetition where the token is present anywhere + (``counts > 0`` or the prefix bitmap), 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 _pack_presence_prefix(counts: torch.Tensor, presence: torch.Tensor) -> torch.Tensor: + packed = torch.zeros( + presence.size(0), + (presence.size(1) + 31) // 32, + dtype=torch.int32, + 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, + packed, + empty, + empty, + prefix_slots, + prefix_tokens, + ) + return packed + + +@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 bitmap 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_triton_matches_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 not a multiple of BLOCK_SIZE (1024), exercising the tail mask. + 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 = _pack_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_triton( + 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_triton_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_triton( + 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 kernel output, then match its fp32-compute -> bf16-store + # boundary. This keeps the tolerance about Triton 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_packed_prefix_boundaries_match_dense_logits_reference() -> None: + vocab = 70 + counts = torch.zeros(1, vocab, dtype=torch.int32, device="cuda") + presence_prefix = torch.zeros(1, (vocab + 31) // 32, dtype=torch.int32, 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_triton( + 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, 3) + torch.testing.assert_close(logits, expected, rtol=1e-4, atol=1e-4) + + +def test_batched_kernel_does_not_latch_pending_token() -> None: + """The batched kernel must not write ``has_previous_token``. + + Every vocab block reads the flag at entry to decide whether to fold the pending + ``new_tokens`` token; a write from the kernel could be observed by a later block + (CTAs are unordered) and spuriously fold ``new_tokens`` into its ``counts`` slice. + Here the flag is False with a stale token in a high vocab block (above ``BLOCK_SIZE``): + nothing may be folded, the flag must stay False, and the logits must be untouched. + """ + vocab = 3000 # > BLOCK_SIZE (1024) -> spans multiple vocab blocks + stale_token = 2500 # lives in vocab block 2, above block 0 + 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_triton( + 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 raw kernel 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 high-block 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, + ) + + +def _apply_handler( + handler: _OccurrencePenaltyHandler, + 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 = _OccurrencePenaltyHandler( + max_num_sequences=3, + device="cuda", + ) + history = [3] + request = _make_handler_request(slot=slot, tokens=history) + handler.prepare_for_new_request(request, slot=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 = _OccurrencePenaltyHandler( + 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) + handler.prepare_for_new_request(first, slot=0) + _apply_handler(handler, first, torch.zeros(1, vocab, device="cuda"), 1, new_tokens) + + second = _make_handler_request(slot=0, tokens=[5]) + handler.prepare_for_new_request(second, slot=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) 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)