diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 7dccd521319c..ab8f258f4b46 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -2621,11 +2621,13 @@ legacy-files: &legacy_files | static-analysis-files: &static_analysis_files | (?x)^( tensorrt_llm/_torch/pyexecutor/sampler/sampler.py | + tensorrt_llm/_torch/pyexecutor/sampler/token_ban.py | tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py | tensorrt_llm/_torch/pyexecutor/sampler/ops/interface.py | tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py | tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py | tests/unittest/_torch/sampler/test_torch_sampler.py | + tests/unittest/_torch/sampler/test_token_ban.py | tests/unittest/_torch/sampler/test_beam_search.py | tests/unittest/_torch/sampler/test_beam_search_util.py | )$ diff --git a/docs/source/features/sampling.md b/docs/source/features/sampling.md index 72ad4a62ad41..29f75abde6ba 100644 --- a/docs/source/features/sampling.md +++ b/docs/source/features/sampling.md @@ -115,6 +115,10 @@ 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. +* 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. + ### Performance The Torch Sampler leverages the optimized sampling kernels provided by diff --git a/pyproject.toml b/pyproject.toml index 49c5584d5680..ac6cd49ab87e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1468,6 +1468,7 @@ module = [ "tests.*", "_torch.*", # tests/unittest/_torch/*, cf. mypy_path "test_torch_sampler", # tests/unittest/_torch/sampler/..., cf. mypy_path + "test_token_ban", # tests/unittest/_torch/sampler/..., cf. mypy_path "test_beam_search", # tests/unittest/_torch/sampler/..., cf. mypy_path "test_beam_search_util", # tests/unittest/_torch/sampler/..., cf. mypy_path ] diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 88513bc326d6..d52e1f086e7e 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -1201,6 +1201,13 @@ def executor_request_to_llm_request( list(word) for word in executor_request.bad_words ] if executor_request.bad_words else None + # No-repeat-ngram size for the TorchSampler path, normalized once here so + # the per-step sampler gate is a plain attribute read. The C++ + # SamplingConfig rejects negative values up front and 0 means disabled + # (same convention as the C++ banRepeatNgram kernel), so falsy == off. + ngram_size = executor_request.sampling_config.no_repeat_ngram_size + llm_request.py_no_repeat_ngram_size = ngram_size if ngram_size else None + llm_request.py_original_end_id = getattr(executor_request, "py_original_end_id", llm_request.py_end_id) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 041cb31f97be..edf6ddd57fc6 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -26,6 +26,7 @@ Any, Callable, Dict, + Final, Generic, List, Optional, @@ -105,6 +106,7 @@ sample_rejected, top_p_decay_active, ) +from .token_ban import OverlappedTokenBanHandler, SynchronousTokenBanHandler, TokenBanHandler if sys.version_info[:2] >= (3, 12): from typing import override @@ -1346,8 +1348,8 @@ def _record_sampler_event( class TorchSampler(Sampler[SampleStateTorch], AsyncWorkerMixin): - DEFAULT_MAX_STOP_WORD_LENGTH = 20 - DEFAULT_MAX_STOP_WORDS = 10 + DEFAULT_MAX_STOP_WORD_LENGTH: Final[int] = 20 + DEFAULT_MAX_STOP_WORDS: Final[int] = 10 SampleState = SampleStateTorch @@ -2462,6 +2464,15 @@ def __init__(self, args: Args): # buffer entries are never consumed. self._top_p_decay_slots: set[int] = set() + # Token-ban handling (bad words, no-repeat ngram). The overlap-aware + # variant is selected once here from whether the overlap scheduler is + # enabled; only it produces the conditional (stale-host) bans. + self._token_ban_handler: TokenBanHandler = ( + OverlappedTokenBanHandler() + if self._track_pending_steps + else SynchronousTokenBanHandler() + ) + # AutoDeploy build creates the sampler in inference mode, # which would disallow in-place mutating of new_tokens. # So, we temporarily exit inference mode. @@ -4736,218 +4747,35 @@ def _update_top_p_decay_after_sample( sampled_slots=sampled_slots_cuda, ) - @staticmethod - @torch.inference_mode() - def _apply_min_length_penalty( - logits: torch.Tensor, - requests: list[LlmRequest], - num_steps_tensor: torch.Tensor, - num_beams_tensor: torch.Tensor, - ) -> None: - """Apply min_length_penalty to logits, mutating ``logits`` in place. - - Args: - logits: The logits to apply min length penalty to - requests: The requests to apply min length penalty to - num_steps_tensor: The number of steps per request (host tensor) - num_beams_tensor: The number of beams per request (host tensor) + def _compute_stale_by_one(self, requests: list[LlmRequest]) -> list[bool] | None: + """Per-request overlap-scheduler stale flags, or None when not applicable. + + Returns a list where entry ``i`` is True when request ``i``'s host token + history lags the device state by exactly one token (the previous step's + token was sampled but not yet written back). Only the single-step, + single-beam overlap case is reconstructible on the device side; under + speculative decoding or beam search the missing history cannot be + recovered, so bans are matched against the lagging host history and may + be enforced one step late (warned once). Returns None when the overlap + scheduler is off, on a draft batch, or when nothing is pending. """ - if not any( - r.py_min_length and (r.max_beam_num_tokens - r.py_orig_prompt_len) < r.py_min_length[0] - for r in requests - ): - return - - # Deferred host conversion: only needed on the (rare) penalty path. - num_steps = num_steps_tensor.tolist() - num_beams = num_beams_tensor.tolist() - - rows: list[int] = [] - cols: list[int] = [] - current_offset = 0 - for index, r in enumerate(requests): - # Advance the offset before any guard below can skip the request: - # every request occupies its logits rows, penalized or not. - req_offset = current_offset - current_offset += num_steps[index] * num_beams[index] - - if not r.py_min_length: - continue - # Use the original end_id (before ignore_eos override) - # so we suppress the real EOS token, not token -1. - end_id = getattr(r, "py_original_end_id", r.py_end_id) - if end_id is None or end_id <= -1: - continue - - for beam_idx in range(num_beams[index]): - for step in range(num_steps[index]): - if (r.get_num_tokens(beam_idx) - r.py_orig_prompt_len) + step < r.py_min_length[ - 0 - ]: - rows.append(req_offset + num_steps[index] * beam_idx + step) - cols.append(end_id) - else: - break - - if rows: - neg_inf = torch.full((), float("-inf"), dtype=logits.dtype, device=logits.device) - row_idx = torch.tensor(rows, dtype=torch.long, pin_memory=prefer_pinned()).to( - logits.device, non_blocking=True - ) - col_idx = torch.tensor(cols, dtype=torch.long, pin_memory=prefer_pinned()).to( - logits.device, non_blocking=True - ) - logits.index_put_((row_idx, col_idx), neg_inf, accumulate=False) - - @staticmethod - @torch.inference_mode() - def _apply_bad_words( - logits: torch.Tensor, - requests: list[LlmRequest], - num_steps: list[int], - num_beams: list[int], - *, - new_tokens_cuda: torch.Tensor | None = None, - stale_by_one: list[bool] | None = None, - ) -> None: - """Inplace ban "bad words" by masking their final token's logit to -inf. - - A single-token word is banned unconditionally; a multi-token word - ``[t0, ..., t_{k-1}]`` bans its final token ``t_{k-1}`` only when the - ``k-1`` most recently generated tokens exactly match the prefix - ``[t0, ..., t_{k-2}]``. - - With the overlap scheduler, ``sample_async`` for step ``i`` runs before - ``update_requests`` for step ``i-1``, so ``r.get_tokens()`` is missing - exactly the previous step's token; that token still lives device-side in - ``new_tokens_cuda``. For requests flagged in ``stale_by_one``, the - prefix match is therefore split: the first ``k-2`` prefix tokens are - matched on the host against the (complete up to there) host context, and - the final prefix token is compared on the GPU against - ``new_tokens_cuda[0, seq_slot, 0]``, without any device-to-host - synchronization. This path only supports ``num_steps == 1`` and - ``num_beams == 1`` (no speculation, no beam search). - - Args: - logits: Flattened ``[total_rows, vocab]`` logits; rows are packed per - request as ``num_steps * num_beams`` consecutive entries, in - beam-major / step-minor order (same layout as - ``_apply_min_length_penalty``). Modified in-place. - requests: The requests, aligned with the packed logits rows. - num_steps: Number of steps per request. - num_beams: Number of beams per request. - new_tokens_cuda: Device buffer holding the previous step's sampled - tokens, shape ``[max_tokens, max_num_sequences, max_beam_width]``. - Required when any entry of ``stale_by_one`` is True. - stale_by_one: Per-request flag; True when the host token list lags - the device state by exactly one token (overlap scheduler). - """ - rows: list[int] = [] - cols: list[int] = [] - # Overlap path: bans whose last prefix token must be compared on-device. - cond_rows: list[int] = [] - cond_cols: list[int] = [] - cond_slots: list[int] = [] - cond_expected: list[int] = [] - current_offset = 0 - for index, r in enumerate(requests): - request_offset = current_offset - # Advance to the next request's rows before any early continue. - current_offset += num_steps[index] * num_beams[index] - - bad_words = getattr(r, "py_bad_words", None) - if not bad_words: - continue - - if stale_by_one is not None and stale_by_one[index]: - assert num_steps[index] == 1 and num_beams[index] == 1, ( - "stale-host bad-words path only supports a single step and beam" - ) - assert r.py_seq_slot is not None - # Host context is missing the previous step's token; the true - # sequence is context + [new_tokens_cuda[0, seq_slot, 0]]. - context = r.get_tokens(0) - for word in bad_words: - k = len(word) - if k == 0: - continue - if k == 1: - # Single-token word: banned unconditionally. - rows.append(request_offset) - cols.append(word[0]) - continue - # True sequence length is len(context) + 1; need >= k - 1. - if len(context) < k - 2: - continue - # Host part: all prefix tokens except the newest one. - if k > 2 and context[-(k - 2) :] != word[: k - 2]: - continue - # Device part: the previous step's token must equal the - # last prefix token; resolved on the GPU below. - cond_rows.append(request_offset) - cond_cols.append(word[-1]) - cond_slots.append(r.py_seq_slot) - cond_expected.append(word[k - 2]) - continue - - for beam_idx in range(num_beams[index]): - # Full token sequence for this beam (prompt + generated), so a - # bad-word prefix ending inside the prompt is also matched. - context = r.get_tokens(beam_idx) - for word in bad_words: - k = len(word) - if k == 0: - continue - if k == 1: - # Single-token word: banned unconditionally. - col = word[0] - elif len(context) >= k - 1 and context[-(k - 1) :] == word[:-1]: - # Multi-token word: ban the final token only when the - # generated suffix matches the word prefix. - col = word[-1] - else: - continue - # Apply to every step row of this beam. - for step in range(num_steps[index]): - rows.append(request_offset + num_steps[index] * beam_idx + step) - cols.append(col) - - if rows: - neg_inf = torch.full((), float("-inf"), dtype=logits.dtype, device=logits.device) - row_idx = torch.tensor(rows, dtype=torch.long, pin_memory=prefer_pinned()).to( - logits.device, non_blocking=True - ) - col_idx = torch.tensor(cols, dtype=torch.long, pin_memory=prefer_pinned()).to( - logits.device, non_blocking=True - ) - logits.index_put_((row_idx, col_idx), neg_inf, accumulate=False) - - if cond_rows: - assert new_tokens_cuda is not None - device = logits.device - cond_row_idx = torch.tensor(cond_rows, dtype=torch.long, pin_memory=prefer_pinned()).to( - device, non_blocking=True - ) - cond_col_idx = torch.tensor(cond_cols, dtype=torch.long, pin_memory=prefer_pinned()).to( - device, non_blocking=True - ) - slot_idx = torch.tensor(cond_slots, dtype=torch.long, pin_memory=prefer_pinned()).to( - device, non_blocking=True - ) - expected = torch.tensor( - cond_expected, dtype=new_tokens_cuda.dtype, pin_memory=prefer_pinned() - ).to(device, non_blocking=True) - # Previous step's token per request (single step, single beam). - prev_tokens = new_tokens_cuda[0].index_select(0, slot_idx)[:, 0] - # -inf where the device-side prefix token matches, 0 otherwise. - # Additive update keeps the op shape-static (boolean-mask indexing - # would force a device-to-host sync). - penalty = torch.where( - prev_tokens == expected, - torch.full((), float("-inf"), dtype=logits.dtype, device=device), - torch.zeros((), dtype=logits.dtype, device=device), - ) - logits.index_put_((cond_row_idx, cond_col_idx), penalty, accumulate=True) + if not self._track_pending_steps or self._is_draft_batch(requests): + return None + pending = [ + self._pending_steps[r.py_seq_slot] if r.py_seq_slot is not None else 0 for r in requests + ] + if not any(pending): + return None + if self.max_tokens == 1 and self.max_beam_width == 1 and max(pending) == 1: + return [p > 0 for p in pending] + logger.warning_once( + "bad_words / no_repeat_ngram_size with the overlap scheduler and " + "speculative decoding or beam search: bans are matched against a " + "host token history that lags the device state and may be enforced " + "inexactly.", + key="bad_words_stale_overlap", + ) + return None @staticmethod def _select_generated_logits( @@ -5254,43 +5082,27 @@ def _process_requests( logits_cuda, sampling_requests, sampling_requests_metadata.req_num_steps ) - self._apply_min_length_penalty( - logits_cuda, - sampling_requests, - sampling_requests_metadata.req_num_steps, - sampling_requests_metadata.req_num_beams, - ) - - if any(getattr(r, "py_bad_words", None) for r in sampling_requests): - stale_by_one: list[bool] | None = None - if self._track_pending_steps and not self._is_draft_batch(sampling_requests): - pending = [ - self._pending_steps[r.py_seq_slot] if r.py_seq_slot is not None else 0 - for r in sampling_requests - ] - if any(pending): - if self.max_tokens == 1 and self.max_beam_width == 1 and max(pending) == 1: - stale_by_one = [p > 0 for p in pending] - else: - # Speculative decoding / beam search under overlap: - # the missing host tokens cannot be reconstructed with - # the single-token device-side check; multi-token bad - # words may be applied one step late. - logger.warning_once( - "bad_words with the overlap scheduler and speculative " - "decoding or beam search: multi-token bad words are " - "matched against a host token history that lags the " - "device state and may be enforced inexactly.", - key="bad_words_stale_overlap", - ) - self._apply_bad_words( - logits_cuda, + 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 + # None when the restriction is disabled for the request. + ngram_sizes = [getattr(r, "py_no_repeat_ngram_size", None) for r in sampling_requests] + has_no_repeat_ngram = any(size is not None for size in ngram_sizes) + if has_min_length or has_bad_words or has_no_repeat_ngram: + # Overlap-scheduler stale flags (per request): True when the host + # token history lags the device by one token. Only the overlap + # handler consumes them; computed here as it needs sampler state. + stale_by_one = self._compute_stale_by_one(sampling_requests) + bans = self._token_ban_handler.generate_ban_list( sampling_requests, sampling_requests_metadata.req_num_steps.tolist(), sampling_requests_metadata.req_num_beams.tolist(), - new_tokens_cuda=new_tokens_cuda, + ngram_sizes, stale_by_one=stale_by_one, ) + self._token_ban_handler.apply_ban_list( + logits_cuda, bans, new_tokens_cuda=new_tokens_cuda + ) # Fast path for greedy sampling if self._can_use_fast_greedy_path(sampling_requests): diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/token_ban.py b/tensorrt_llm/_torch/pyexecutor/sampler/token_ban.py new file mode 100644 index 000000000000..c4a44112e8f1 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/sampler/token_ban.py @@ -0,0 +1,673 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Token-ban handling for the TorchSampler. + +A *token ban* masks specific logits to ``-inf`` so certain tokens cannot be +sampled. Several features reduce to this: bad words, no-repeat n-gram, and the +min-length EOS suppression. + +All bans are accumulated on the host into a single :class:`TokenBans` and then +applied to the logits in one batched step per category. Two handler variants +exist, selected once at construction from whether the overlap scheduler is +enabled: + +* :class:`SynchronousTokenBanHandler` -- overlap disabled. The host token + history is always complete, so every ban is unconditional. +* :class:`OverlappedTokenBanHandler` -- overlap enabled. The host history may + lag the device state by one token, so a batch mixes "fresh" requests (history + complete) with "stale" ones (missing the previous step's token, still on the + device); the latter produce conditional bans resolved on the device at apply + time without a device-to-host sync. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections import defaultdict +from collections.abc import Callable, Iterable +from dataclasses import dataclass, field +from weakref import WeakKeyDictionary + +import torch + +from tensorrt_llm._utils import prefer_pinned + +from ..llm_request import LlmRequest + +# A suffix-ban rule ``(prefix, col)``: ban token ``col`` when the sequence ends +# with ``prefix`` (an empty prefix bans unconditionally). +BanRule = tuple[list[int], int] +RuleGenerator = Callable[[int, LlmRequest, Callable[[], list[int]]], Iterable[BanRule]] + + +def _memoizing_tokens_getter(request: LlmRequest, beam_idx: int) -> Callable[[], list[int]]: + """Return a thunk for ``request.get_tokens(beam_idx)`` that fetches once. + + ``get_tokens`` copies the whole token history through the bindings (~40us + at 2048 tokens), so rule generators receive this thunk and only pay for + the copy if they actually need the history. + """ + tokens: list[int] | None = None + + def get() -> list[int]: + nonlocal tokens + if tokens is None: + tokens = request.get_tokens(beam_idx) + return tokens + + return get + + +@dataclass(kw_only=True) +class TokenBans: + """Accumulated token bans, applied by ``TokenBanHandler.apply_ban_list``. + + Host-side accumulator; nothing here lives on the GPU. ``apply_ban_list`` + turns the populated lists into tensors and writes ``-inf`` into the logits, + one batched transfer and ``index_put_`` per category. + """ + + # Unconditional bans: logits[row, col] = -inf (fresh requests + min_length). + rows: list[int] = field(default_factory=list) + cols: list[int] = field(default_factory=list) + # Conditional bans (overlap/stale): ban only where the device-side previous + # token new_tokens_cuda[0, slot, 0] == expected. + cond_rows: list[int] = field(default_factory=list) + cond_cols: list[int] = field(default_factory=list) + cond_slots: list[int] = field(default_factory=list) + cond_expected: list[int] = field(default_factory=list) + # Conditional special case: the banned column IS the device token itself + # (overlap + no_repeat_ngram n == 1). Only OverlappedTokenBanHandler fills + # these. + dev_rows: list[int] = field(default_factory=list) + dev_slots: list[int] = field(default_factory=list) + + +@dataclass +class _NgramIndexCache: + """Incrementally maintained n-gram index for a single request (single beam). + + A request's ngram size is fixed for its lifetime, so the index is only ever + grown, never rebuilt for a different ``n``; the size is therefore not stored. + """ + + tokens: list[int] + """Host-side mirror of the request's token history.""" + indexed_up_to: int + """How far into ``tokens`` the indices have been built, so each window is + indexed exactly once.""" + fresh_idx: defaultdict[tuple[int, ...], set[int]] + """Maps each ``(n - 1)``-gram to the tokens that followed it.""" + stale_idx: defaultdict[tuple[int, ...], set[tuple[int, int]]] + """Populated only for ``n > 2``: maps each ``(n - 2)``-gram to its + ``(next, next-but-one)`` token pairs.""" + + +class TokenBanHandler(ABC): + """Base class for token-ban handling: state + shared collect/apply helpers. + + Subclasses implement ``generate_ban_list`` (build a :class:`TokenBans`) and + ``apply_ban_list`` (write it into the logits). Shared feature logic (bad + words, no-repeat ngram, min-length EOS suppression, the suffix matcher, and + the unconditional flush) lives here. + """ + + def __init__(self) -> None: + # Per-request incrementally maintained n-gram indices, keyed weakly on + # the request so entries are reclaimed when the request is freed. + self._ngram_index_caches: WeakKeyDictionary[LlmRequest, _NgramIndexCache] = ( + WeakKeyDictionary() + ) + + # ---- entry points (subclass-specific) -------------------------------- + + @abstractmethod + def generate_ban_list( + self, + requests: list[LlmRequest], + num_steps: list[int], + num_beams: list[int], + ngram_sizes: list[int | None], + *, + stale_by_one: list[bool] | None = None, + ) -> TokenBans: + """Collect all enabled feature bans into a fresh :class:`TokenBans`. + + ``stale_by_one`` is the per-request overlap-scheduler flag; it is only + meaningful for :class:`OverlappedTokenBanHandler` and ignored otherwise. + """ + + @abstractmethod + def apply_ban_list( + self, + logits: torch.Tensor, + bans: TokenBans, + *, + new_tokens_cuda: torch.Tensor | None = None, + ) -> None: + """Write the accumulated bans into ``logits`` in-place.""" + + # ---- feature rule sources (shared) ----------------------------------- + + def _add_bad_words_bans( + self, + bans: TokenBans, + requests: list[LlmRequest], + num_steps: list[int], + num_beams: list[int], + *, + stale_by_one: list[bool] | None, + ) -> None: + """Add "bad words" bans. + + A single-token word is banned unconditionally; a multi-token word + ``[t0, ..., t_{k-1}]`` bans its final token ``t_{k-1}`` only when the + ``k-1`` most recently generated tokens match the prefix + ``[t0, ..., t_{k-2}]``. + """ + + def fresh_rules( + index: int, r: LlmRequest, get_context: Callable[[], list[int]] + ) -> list[BanRule]: + # A bad word is the suffix rule "after word[:-1], ban word[-1]"; + # matching is left to the driver. + return [ + (word[:-1], word[-1]) for word in getattr(r, "py_bad_words", None) or () if word + ] + + def stale_rules( + index: int, r: LlmRequest, get_context: Callable[[], list[int]] + ) -> list[BanRule]: + # Stale contract: check the first k - 2 prefix tokens here; the + # last one is compared against the device token d at flush time. + rules: list[BanRule] = [] + context: list[int] | None = None + for word in getattr(r, "py_bad_words", None) or (): + k = len(word) + if k == 0: + continue + if k == 1: + rules.append(([], word[0])) + continue + if context is None: + context = get_context() + # True sequence length is len(context) + 1; need >= k - 1. + if len(context) < k - 2: + continue + if k > 2 and context[-(k - 2) :] != word[: k - 2]: + continue + rules.append((word[:-1], word[-1])) + return rules + + self._add_suffix_rule_bans( + bans, + requests, + num_steps, + num_beams, + active=lambda index, r: bool(getattr(r, "py_bad_words", None)), + fresh_rules=fresh_rules, + stale_rules=stale_rules, + stale_by_one=stale_by_one, + ) + + def _add_min_length_bans( + self, + bans: TokenBans, + requests: list[LlmRequest], + num_steps: list[int], + num_beams: list[int], + ) -> None: + """Add bans that suppress EOS until a request reaches its min length. + + For a request with ``py_min_length``, the end-of-sequence token is + banned on every step whose generated length is still below the minimum. + This is an unconditional ban that does not depend on the token history, + so it has no stale-host variant and is identical under both handlers. + """ + current_offset = 0 + for index, r in enumerate(requests): + # Advance the offset before any guard below can skip the request: + # every request occupies its logits rows, penalized or not. + request_offset = current_offset + current_offset += num_steps[index] * num_beams[index] + + if not r.py_min_length: + continue + # Use the original end_id (before ignore_eos override) so we + # suppress the real EOS token, not token -1. + end_id = getattr(r, "py_original_end_id", r.py_end_id) + if end_id is None or end_id <= -1: + continue + + for beam_idx in range(num_beams[index]): + for step in range(num_steps[index]): + if (r.get_num_tokens(beam_idx) - r.py_orig_prompt_len) + step < r.py_min_length[ + 0 + ]: + bans.rows.append(request_offset + num_steps[index] * beam_idx + step) + bans.cols.append(end_id) + else: + break + + def _add_no_repeat_ngram_bans( + self, + bans: TokenBans, + requests: list[LlmRequest], + num_steps: list[int], + num_beams: list[int], + ngram_sizes: list[int | None], + *, + stale_by_one: list[bool] | None, + ) -> None: + """Add bans for tokens that would repeat an existing n-gram. + + For no-repeat-ngram size ``n``, the last ``n - 1`` tokens of the + sequence (prompt + generated) form the current prefix; the final token + of every existing n-gram whose first ``n - 1`` tokens equal that prefix + is masked to ``-inf`` (same semantics as the C++ ``banRepeatNgram`` + kernel and HF's ``NoRepeatNGramLogitsProcessor``). ``n == 1`` bans every + token already present. The restriction is per-beam. + + On the stale-host overlap path the rules cover the true sequence + ``tokens + [d]``: the window ending at ``d`` reduces to a host-known + column, and ``n == 1`` additionally bans ``d`` itself (``stale_extra``). + + ``ngram_sizes`` holds the effective per-request size (None or a + non-positive value disables the restriction). + """ + + def fresh_rules( + index: int, r: LlmRequest, get_context: Callable[[], list[int]] + ) -> list[BanRule]: + # Every existing n-gram is the suffix rule "after its first n - 1 + # tokens, ban its last token". + n = ngram_sizes[index] + if not n: # None or non-positive: restriction disabled + return [] + if num_beams[index] == 1: + # Single beam: O(1) index lookup, rules come out pre-matched. + tokens, fresh_idx, _ = self._extend_ngram_index(r, n) + if len(tokens) < n: + return [] + key = tuple(tokens[-(n - 1) :]) if n > 1 else () + return [([], col) for col in fresh_idx.get(key, ())] + # Multi-beam: histories diverge, so scan this beam's context. + context = get_context() + length = len(context) + if length < n: + return [] + if n == 1: + return [([], col) for col in set(context)] + return [(context[i : i + n - 1], context[i + n - 1]) for i in range(length - n + 1)] + + def stale_rules( + index: int, r: LlmRequest, get_context: Callable[[], list[int]] + ) -> list[BanRule]: + # Stale contract: rules are already host-matched; only the device + # comparison of d against prefix[-1] remains. + n = ngram_sizes[index] + if not n: # None or non-positive: restriction disabled + return [] + tokens, fresh_idx, stale_idx = self._extend_ngram_index(r, n) + if n == 1: + # Host tokens; d itself is banned via stale_extra below. + return [([], col) for col in fresh_idx.get((), ())] + if n == 2: + # Empty (n - 2)-gram head: every indexed bigram is a candidate; + # reconstruct the pairs from fresh_idx. + pairs: Iterable[tuple[int, int]] = ( + (key[0], col) for key, cols in fresh_idx.items() for col in cols + ) + else: + pairs = stale_idx.get(tuple(tokens[-(n - 2) :]), ()) + rules: list[BanRule] = [([expected], col) for expected, col in pairs] + # The window ending at d bans d itself; the match condition forces + # d == tokens[-1] (host-known). Check its host-side part here. + if len(tokens) >= n - 1 and (n == 2 or tokens[-(n - 1) : -1] == tokens[-(n - 2) :]): + rules.append(([tokens[-1]], tokens[-1])) + return rules + + def stale_extra(bans: TokenBans, request_offset: int, index: int, r: LlmRequest) -> None: + if ngram_sizes[index] == 1: + # n == 1 also bans the device-only token, via a gathered column. + assert r.py_seq_slot is not None + bans.dev_rows.append(request_offset) + bans.dev_slots.append(r.py_seq_slot) + + self._add_suffix_rule_bans( + bans, + requests, + num_steps, + num_beams, + active=lambda index, r: bool(ngram_sizes[index]), + fresh_rules=fresh_rules, + stale_rules=stale_rules, + stale_extra=stale_extra, + stale_by_one=stale_by_one, + ) + + def _extend_ngram_index( + self, request: LlmRequest, n: int + ) -> tuple[ + list[int], dict[tuple[int, ...], set[int]], dict[tuple[int, ...], set[tuple[int, int]]] + ]: + """Incrementally maintained n-gram index for a single-beam request. + + Returns ``(tokens, fresh_idx, stale_idx)``: a host-side mirror of the + token history, a map from each ``(n - 1)``-gram to the tokens that + followed it, and (for ``n > 2`` only) a map from each ``(n - 2)``-gram + to its ``(next, next-but-one)`` token pairs. The mirror grows via the + scalar ``get_num_tokens`` / ``get_last_tokens`` accessors — a full + ``get_tokens()`` copy costs ~40us at 2048 tokens — and each window is + indexed once, so the steady-state per-step cost is O(new tokens). + Rebuilt when the history shrinks (speculative rollback); single-beam + only. + """ + num_tokens = request.get_num_tokens(0) + cache = self._ngram_index_caches.get(request) + # A request's ngram size is fixed for its lifetime and the cache is + # keyed on the request, so the cached index is always for this ``n``. + if cache is not None and len(cache.tokens) <= num_tokens: + tokens, indexed_up_to = cache.tokens, cache.indexed_up_to + fresh_idx, stale_idx = cache.fresh_idx, cache.stale_idx + if len(tokens) == num_tokens - 1: + # Common decode step: exactly one new token since last call. + tokens.append(request.get_last_tokens(0)) + elif len(tokens) < num_tokens: + # Several tokens landed at once (e.g. accepted draft tokens). + tokens = request.get_tokens(0) + else: + # First use or shrunken history (speculative rollback). + tokens, indexed_up_to = request.get_tokens(0), 0 + fresh_idx, stale_idx = defaultdict(set), defaultdict(set) + end = max(len(tokens) - n + 1, 0) + for i in range(indexed_up_to, end): + # defaultdict(set) creates the entry on first access, avoiding an + # empty set constructed (and discarded) on every already-seen key. + fresh_idx[tuple(tokens[i : i + n - 1])].add(tokens[i + n - 1]) + if n > 2: + stale_idx[tuple(tokens[i : i + n - 2])].add((tokens[i + n - 2], tokens[i + n - 1])) + self._ngram_index_caches[request] = _NgramIndexCache( + tokens=tokens, indexed_up_to=end, fresh_idx=fresh_idx, stale_idx=stale_idx + ) + return tokens, fresh_idx, stale_idx + + # ---- shared suffix matcher + packed-row bookkeeping ------------------ + + def _add_suffix_rule_bans( + self, + bans: TokenBans, + requests: list[LlmRequest], + num_steps: list[int], + num_beams: list[int], + *, + active: Callable[[int, LlmRequest], bool], + fresh_rules: RuleGenerator, + stale_rules: RuleGenerator, + stale_extra: Callable[[TokenBans, int, int, LlmRequest], None] | None = None, + stale_by_one: list[bool] | None = None, + ) -> None: + """Shared matcher for suffix-rule token-ban features. + + A rule ``(prefix, col)`` bans token ``col`` when the sequence (prompt + + generated) ends with ``prefix``; an empty prefix bans unconditionally. + Both bad words and no-repeat ngram reduce to sets of such rules; this + driver owns the packed-row bookkeeping and the suffix matching. + + Stale-host overlap path: with the overlap scheduler, ``sample_async`` + for step ``i`` runs before ``update_requests`` for step ``i - 1``, so + the host token list misses exactly the previous step's token ``d``, + which still lives device-side in ``new_tokens_cuda[0, seq_slot, 0]``. + For requests flagged in ``stale_by_one``, a rule's suffix match against + the true sequence ``context + [d]`` is split into a host-side check of + ``prefix[:-1]`` and a device-side comparison of ``d`` against + ``prefix[-1]``, resolved at flush time without any device-to-host + synchronization. This path only supports ``num_steps == 1`` and + ``num_beams == 1`` (no speculation, no beam search). + + Args: + bans: Accumulator the collected bans are appended to; row indices + refer to the flattened ``[total_rows, vocab]`` logits layout + (``num_steps * num_beams`` consecutive rows per request, in + beam-major / step-minor order). + requests: The requests, aligned with the packed logits rows. + num_steps: Number of steps per request. + num_beams: Number of beams per request. + active: ``(index, request) -> bool``; requests where this is False + are skipped entirely. + fresh_rules: ``(index, request, get_context) -> iterable`` of + ``(prefix, col)`` rules; the driver matches them against the + beam context. ``get_context()`` fetches (and memoizes) the + beam's token history — a full copy through the bindings — so + generators should call it only when needed. + stale_rules: Like ``fresh_rules``, but for the true sequence + ``context + [d]`` and already host-matched by the generator: + the driver only compares ``d`` against ``prefix[-1]`` on the + device (empty prefix = unconditional). Prefixes and cols must + consist of host-known tokens. + stale_extra: Optional hook ``(bans, request_offset, index, request)`` + invoked once per stale request for bans the rule form cannot + express (e.g. banning the device token itself). + stale_by_one: Per-request flag; True when the host token list lags + the device state by exactly one token (overlap scheduler). + These bans need ``new_tokens_cuda`` when flushed. + """ + current_offset = 0 + for index, r in enumerate(requests): + request_offset = current_offset + # Advance to the next request's rows before any early continue. + current_offset += num_steps[index] * num_beams[index] + + if not active(index, r): + continue + + if stale_by_one is not None and stale_by_one[index]: + assert num_steps[index] == 1 and num_beams[index] == 1, ( + "stale-host token-ban path only supports a single step and beam" + ) + assert r.py_seq_slot is not None + # Rules arrive host-matched; the device comparison of d against + # prefix[-1] is resolved at flush time. + seen_uncond: set[int] = set() + seen_cond: set[tuple[int, int]] = set() + for prefix, col in stale_rules(index, r, _memoizing_tokens_getter(r, 0)): + if not prefix: + if col not in seen_uncond: + seen_uncond.add(col) + bans.rows.append(request_offset) + bans.cols.append(col) + continue + key = (prefix[-1], col) + if key in seen_cond: + continue + seen_cond.add(key) + bans.cond_rows.append(request_offset) + bans.cond_cols.append(col) + bans.cond_slots.append(r.py_seq_slot) + bans.cond_expected.append(prefix[-1]) + if stale_extra is not None: + stale_extra(bans, request_offset, index, r) + continue + + for beam_idx in range(num_beams[index]): + # The context (prompt + generated) is fetched at most once, + # and only if a rule with a non-empty prefix needs it. + get_context = _memoizing_tokens_getter(r, beam_idx) + seen_cols: set[int] = set() + for prefix, col in fresh_rules(index, r, get_context): + if col in seen_cols: + continue + k = len(prefix) + if k > 0: + context = get_context() + if len(context) < k or context[-k:] != prefix: + continue + seen_cols.add(col) + # Apply to every step row of this beam. With speculation + # the banned set for later steps is approximated from the + # host context. + for step in range(num_steps[index]): + bans.rows.append(request_offset + num_steps[index] * beam_idx + step) + bans.cols.append(col) + + # ---- shared unconditional apply (both variants) ---------------------- + + @staticmethod + @torch.inference_mode() + def _apply_unconditional_bans(logits: torch.Tensor, bans: TokenBans) -> None: + """Write the unconditional bans (``bans.rows`` / ``bans.cols``) to -inf. + + Rows and cols are packed into one ``[2, N]`` tensor so the host-to-device + copy is a single transfer; the two index rows are split on the device. + """ + if not bans.rows: + return + neg_inf = torch.full((), float("-inf"), dtype=logits.dtype, device=logits.device) + rowcol_idx = torch.tensor( + [bans.rows, bans.cols], dtype=torch.long, pin_memory=prefer_pinned() + ).to(logits.device, non_blocking=True) + logits.index_put_((rowcol_idx[0], rowcol_idx[1]), neg_inf, accumulate=False) + + +class SynchronousTokenBanHandler(TokenBanHandler): + """Token bans without the overlap scheduler: the host history is complete, + so every request is "fresh" and all bans are unconditional.""" + + def generate_ban_list( + self, + requests: list[LlmRequest], + num_steps: list[int], + num_beams: list[int], + ngram_sizes: list[int | None], + *, + stale_by_one: list[bool] | None = None, + ) -> TokenBans: + assert stale_by_one is None, "synchronous handler has no stale requests" + bans = TokenBans() + if any(getattr(r, "py_min_length", None) for r in requests): + self._add_min_length_bans(bans, requests, num_steps, num_beams) + if any(getattr(r, "py_bad_words", None) for r in requests): + self._add_bad_words_bans(bans, requests, num_steps, num_beams, stale_by_one=None) + if any(size is not None for size in ngram_sizes): + self._add_no_repeat_ngram_bans( + bans, requests, num_steps, num_beams, ngram_sizes, stale_by_one=None + ) + # Without overlap there are no stale requests, hence no conditional bans. + assert not bans.cond_rows and not bans.dev_rows + return bans + + def apply_ban_list( + self, + logits: torch.Tensor, + bans: TokenBans, + *, + new_tokens_cuda: torch.Tensor | None = None, + ) -> None: + self._apply_unconditional_bans(logits, bans) + + +class OverlappedTokenBanHandler(TokenBanHandler): + """Token bans with the overlap scheduler: the host history can lag the + device by one token, so a batch mixes fresh and stale requests. Stale + requests produce conditional bans resolved on the device at apply time. + + The per-request ``stale_by_one`` flags are computed by the sampler (which + owns the pending-step counter and the draft-batch / step / beam-width + context) and passed into ``generate_ban_list``. + """ + + def generate_ban_list( + self, + requests: list[LlmRequest], + num_steps: list[int], + num_beams: list[int], + ngram_sizes: list[int | None], + *, + stale_by_one: list[bool] | None = None, + ) -> TokenBans: + bans = TokenBans() + if any(getattr(r, "py_min_length", None) for r in requests): + # min_length has no stale variant; identical under both handlers. + self._add_min_length_bans(bans, requests, num_steps, num_beams) + if any(getattr(r, "py_bad_words", None) for r in requests): + self._add_bad_words_bans( + bans, requests, num_steps, num_beams, stale_by_one=stale_by_one + ) + if any(size is not None for size in ngram_sizes): + self._add_no_repeat_ngram_bans( + bans, requests, num_steps, num_beams, ngram_sizes, stale_by_one=stale_by_one + ) + return bans + + def apply_ban_list( + self, + logits: torch.Tensor, + bans: TokenBans, + *, + new_tokens_cuda: torch.Tensor | None = None, + ) -> None: + self._apply_unconditional_bans(logits, bans) # fresh requests (+ min_length) + self._apply_conditional_bans(logits, bans, new_tokens_cuda) # stale requests + + @staticmethod + @torch.inference_mode() + def _apply_conditional_bans( + logits: torch.Tensor, + bans: TokenBans, + new_tokens_cuda: torch.Tensor | None, + ) -> None: + """Write the overlap/stale bans, resolving the device-token comparison. + + ``dev_*``: the banned column is the device token ``d`` itself. + ``cond_*``: ban only where ``d == expected``; the additive + ``torch.where`` update keeps the op shape-static (boolean-mask indexing + would force a device-to-host sync). + """ + neg_inf = torch.full((), float("-inf"), dtype=logits.dtype, device=logits.device) + device = logits.device + + if bans.dev_rows: + assert new_tokens_cuda is not None + dev_row_idx = torch.tensor( + bans.dev_rows, dtype=torch.long, pin_memory=prefer_pinned() + ).to(device, non_blocking=True) + dev_slot_idx = torch.tensor( + bans.dev_slots, dtype=torch.long, pin_memory=prefer_pinned() + ).to(device, non_blocking=True) + prev_tokens = new_tokens_cuda[0].index_select(0, dev_slot_idx)[:, 0].to(torch.long) + logits.index_put_((dev_row_idx, prev_tokens), neg_inf, accumulate=False) + + if bans.cond_rows: + assert new_tokens_cuda is not None + cond_row_idx = torch.tensor( + bans.cond_rows, dtype=torch.long, pin_memory=prefer_pinned() + ).to(device, non_blocking=True) + cond_col_idx = torch.tensor( + bans.cond_cols, dtype=torch.long, pin_memory=prefer_pinned() + ).to(device, non_blocking=True) + slot_idx = torch.tensor( + bans.cond_slots, dtype=torch.long, pin_memory=prefer_pinned() + ).to(device, non_blocking=True) + expected = torch.tensor( + bans.cond_expected, dtype=new_tokens_cuda.dtype, pin_memory=prefer_pinned() + ).to(device, non_blocking=True) + # Previous step's token per request (single step, single beam). + prev_tokens = new_tokens_cuda[0].index_select(0, slot_idx)[:, 0] + penalty = torch.where( + prev_tokens == expected, + neg_inf, + torch.zeros((), dtype=logits.dtype, device=device), + ) + logits.index_put_((cond_row_idx, cond_col_idx), penalty, accumulate=True) diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index e7fed5e2ced0..a725a201f043 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -227,7 +227,7 @@ class SamplingParams: prompt_ignore_length (int, optional): Controls how many tokens to ignore from the prompt for presence and frequency penalties. Values <= 0 have no effect. Values > input (prompt) length will be clamped. None means using C++ runtime default 0. Defaults to None. length_penalty (float, optional): Controls how to penalize longer sequences in beam search. None means using C++ runtime default 0.f. Defaults to None. early_stopping (int, optional): Controls whether the generation process finishes once beamWidth sentences are generated (ends with end_token). None means using C++ runtime default 1. Defaults to None. - no_repeat_ngram_size (int, optional): Controls how many repeat ngram size are acceptable. None means using C++ runtime default 1 << 30. Defaults to None. + no_repeat_ngram_size (int, optional): Forbids repeating any n-gram of this size: a token is excluded from sampling if it would recreate an n-gram that already occurs in the sequence (prompt included). None or 0 disables the restriction. Defaults to None. min_p (float, optional): scale the most likely token to determine the minimum token probability. None means using C++ runtime default 0.0. Defaults to None. beam_width_array (List[int], optional): The array of beam width using in Variable-Beam-Width-Search. Defaults to None. diff --git a/tests/unittest/_torch/sampler/test_token_ban.py b/tests/unittest/_torch/sampler/test_token_ban.py new file mode 100644 index 000000000000..ab1de592df6a --- /dev/null +++ b/tests/unittest/_torch/sampler/test_token_ban.py @@ -0,0 +1,520 @@ +# Copyright (c) 2025-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. + + +"""Unit tests for the token-ban handlers (bad words, no-repeat ngram, +min-length EOS suppression) in +tensorrt_llm/_torch/pyexecutor/sampler/token_ban.py.""" + +from typing import cast + +import pytest +import torch + +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest +from tensorrt_llm._torch.pyexecutor.sampler.token_ban import ( + OverlappedTokenBanHandler, + SynchronousTokenBanHandler, +) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +class TestApplyBadWords: + """Unit tests for the bad-words path of TokenBanHandler. + + Single-token words are banned unconditionally; multi-token words ban their + final token only when the token suffix (prompt + generated) matches the + word prefix. + """ + + VOCAB = 16 + + class MockLlmRequest: + """Minimal stub exposing the attributes the bad-words path reads.""" + + def __init__(self, tokens, *, bad_words=None, prompt_len=0, seq_slot=None): + # get_tokens(beam) returns the full token sequence (prompt + + # generated); py_orig_prompt_len marks where generation starts. + self.py_orig_prompt_len = prompt_len + self._tokens = list(tokens) + self.py_bad_words = bad_words + self.py_seq_slot = seq_slot + + def get_tokens(self, beam_idx): + return self._tokens + + def _run(self, requests, num_steps, num_beams): + total_rows = sum(s * b for s, b in zip(num_steps, num_beams)) + logits = torch.zeros(total_rows, self.VOCAB, device="cuda") + ngram_sizes: list[int | None] = [None] * len(requests) + handler = SynchronousTokenBanHandler() + bans = handler.generate_ban_list( + cast(list[LlmRequest], requests), num_steps, num_beams, ngram_sizes + ) + handler.apply_ban_list(logits, bans) + return logits + + @staticmethod + def _banned_cols(logits_row): + return set(torch.nonzero(torch.isinf(logits_row)).flatten().tolist()) + + def test_single_token_unconditional(self): + req = self.MockLlmRequest(tokens=[3, 4], bad_words=[[7]]) + logits = self._run([req], num_steps=[1], num_beams=[1]) + assert self._banned_cols(logits[0]) == {7} + + def test_multi_token_prefix_hit(self): + # tokens end with [9]; word [9, 2] -> ban token 2. + req = self.MockLlmRequest(tokens=[1, 9], bad_words=[[9, 2]]) + logits = self._run([req], num_steps=[1], num_beams=[1]) + assert self._banned_cols(logits[0]) == {2} + + def test_multi_token_prefix_miss(self): + # tokens end with [3]; word [9, 2] prefix [9] does not match. + req = self.MockLlmRequest(tokens=[1, 3], bad_words=[[9, 2]]) + logits = self._run([req], num_steps=[1], num_beams=[1]) + assert self._banned_cols(logits[0]) == set() + + def test_multi_token_prefix_in_prompt(self): + # The prefix [9] lies in the prompt (nothing generated yet); the word + # [9, 2] must still ban token 2. + req = self.MockLlmRequest(tokens=[1, 9], bad_words=[[9, 2]], prompt_len=2) + logits = self._run([req], num_steps=[1], num_beams=[1]) + assert self._banned_cols(logits[0]) == {2} + + # --- Overlap-scheduler (stale-host) path ------------------------------- + # + # With the overlap scheduler the host token list lags the device state by + # one token; the newest token is read from new_tokens_cuda[0, seq_slot, 0] + # on the GPU. These tests drive the overlap handler with stale_by_one set. + + NUM_SLOTS = 4 + + def _new_tokens_cuda(self, slot_tokens): + buf = torch.full((1, self.NUM_SLOTS, 1), -1, dtype=torch.int32, device="cuda") + for slot, tok in slot_tokens.items(): + buf[0, slot, 0] = tok + return buf + + def _run_stale(self, requests, stale_by_one, slot_tokens): + total_rows = len(requests) + logits = torch.zeros(total_rows, self.VOCAB, device="cuda") + ngram_sizes: list[int | None] = [None] * len(requests) + num_steps = [1] * len(requests) + num_beams = [1] * len(requests) + handler = OverlappedTokenBanHandler() + bans = handler.generate_ban_list( + cast(list[LlmRequest], requests), + num_steps, + num_beams, + ngram_sizes, + stale_by_one=stale_by_one, + ) + handler.apply_ban_list(logits, bans, new_tokens_cuda=self._new_tokens_cuda(slot_tokens)) + return logits + + def test_stale_two_token_device_hit(self): + # Host context [1]; device holds the pending token 9; word [9, 2] + # completes its prefix on the device side -> ban token 2. + req = self.MockLlmRequest(tokens=[1], bad_words=[[9, 2]], seq_slot=0) + logits = self._run_stale([req], [True], {0: 9}) + assert self._banned_cols(logits[0]) == {2} + + def test_stale_two_token_device_miss(self): + # Device token is 3, not the required prefix 9 -> nothing banned. + req = self.MockLlmRequest(tokens=[1], bad_words=[[9, 2]], seq_slot=0) + logits = self._run_stale([req], [True], {0: 3}) + assert self._banned_cols(logits[0]) == set() + + def test_stale_three_token_host_and_device_hit(self): + # Word [5, 9, 2]: host suffix must be [5], device token must be 9. + req = self.MockLlmRequest(tokens=[1, 5], bad_words=[[5, 9, 2]], seq_slot=1) + logits = self._run_stale([req], [True], {1: 9}) + assert self._banned_cols(logits[0]) == {2} + + def test_stale_three_token_host_miss(self): + # Host suffix [3] does not match the word prefix [5]; the device token + # matching is irrelevant -> nothing banned. + req = self.MockLlmRequest(tokens=[1, 3], bad_words=[[5, 9, 2]], seq_slot=1) + logits = self._run_stale([req], [True], {1: 9}) + assert self._banned_cols(logits[0]) == set() + + def test_stale_single_token_unconditional(self): + # Single-token words are banned regardless of the device token. + req = self.MockLlmRequest(tokens=[1], bad_words=[[7]], seq_slot=0) + logits = self._run_stale([req], [True], {0: 3}) + assert self._banned_cols(logits[0]) == {7} + + def test_stale_and_fresh_requests_mixed(self): + # Request 0 (stale) matches via the device token; request 1 (fresh) + # matches via the host context, as on the regular path. + stale_req = self.MockLlmRequest(tokens=[1], bad_words=[[9, 2]], seq_slot=0) + fresh_req = self.MockLlmRequest(tokens=[1, 9], bad_words=[[9, 4]], seq_slot=1) + logits = self._run_stale([stale_req, fresh_req], [True, False], {0: 9}) + assert self._banned_cols(logits[0]) == {2} + assert self._banned_cols(logits[1]) == {4} + + def test_stale_conditional_and_unconditional_same_cell(self): + # An unconditional single-token ban and a device-conditional ban on the + # same logit must combine to -inf (no NaN from -inf + -inf). + req = self.MockLlmRequest(tokens=[1], bad_words=[[2], [9, 2]], seq_slot=0) + logits = self._run_stale([req], [True], {0: 9}) + assert self._banned_cols(logits[0]) == {2} + assert not torch.isnan(logits).any() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +class TestApplyNoRepeatNgram: + """Unit tests for the no-repeat-ngram path of TokenBanHandler. + + With ngram size n, the last n-1 tokens of the sequence form the current + prefix; the final token of every existing n-gram whose first n-1 tokens + match that prefix is banned, so no n-gram is generated twice. + """ + + VOCAB = 16 + + class MockLlmRequest: + """Minimal stub exposing the attributes the no-repeat-ngram path reads.""" + + def __init__(self, tokens, *, seq_slot=None): + self._tokens = list(tokens) + self.py_seq_slot = seq_slot + + def get_tokens(self, beam_idx): + # Copy, like the real binding: the sampler's mirror must not alias. + return list(self._tokens) + + def get_num_tokens(self, beam_idx): + return len(self._tokens) + + def get_last_tokens(self, beam_idx): + return self._tokens[-1] + + class MultiBeamMockLlmRequest: + """Stub with a distinct token history per beam for multi-beam tests.""" + + def __init__(self, beam_tokens, *, seq_slot=None): + self._beam_tokens = [list(t) for t in beam_tokens] + self.py_seq_slot = seq_slot + + def get_tokens(self, beam_idx): + return list(self._beam_tokens[beam_idx]) + + def get_num_tokens(self, beam_idx): + return len(self._beam_tokens[beam_idx]) + + def get_last_tokens(self, beam_idx): + return self._beam_tokens[beam_idx][-1] + + @staticmethod + def _distinct_logits(total_rows, vocab): + # Distinct non-zero values per cell, so a comparison against the input + # can detect any logit the sampler touched but should not have. + return (torch.arange(total_rows * vocab, device="cuda", dtype=torch.float) + 1.0).reshape( + total_rows, vocab + ) + + def _run(self, requests, ngram_sizes, num_steps=None, num_beams=None, handler=None): + num_steps = num_steps or [1] * len(requests) + num_beams = num_beams or [1] * len(requests) + total_rows = sum(s * b for s, b in zip(num_steps, num_beams, strict=True)) + logits = self._distinct_logits(total_rows, self.VOCAB) + before = logits.clone() + handler = handler or SynchronousTokenBanHandler() + bans = handler.generate_ban_list( + cast(list[LlmRequest], requests), num_steps, num_beams, ngram_sizes + ) + handler.apply_ban_list(logits, bans) + self._assert_only_banned_changed(logits, before) + return logits + + @staticmethod + def _banned_cols(logits_row): + return set(torch.nonzero(torch.isinf(logits_row)).flatten().tolist()) + + @classmethod + def _assert_only_banned_changed(cls, logits, before): + # Every non-banned cell must still hold its original (non-zero) value; + # banned cells become -inf. + banned = torch.isinf(logits) + assert torch.equal(logits[~banned], before[~banned]), ( + "a logit outside the banned set was modified" + ) + + def test_bigram_repeat_banned(self): + # Sequence [1, 2, 3, 1, 2] with n=2: prefix [2]; existing bigram + # (2, 3) -> ban 3. The trailing (1, 2) is the prefix itself. + req = self.MockLlmRequest(tokens=[1, 2, 3, 1, 2]) + logits = self._run([req], [2]) + assert self._banned_cols(logits[0]) == {3} + + def test_trigram_repeat_banned(self): + # Sequence [1, 2, 3, 4, 1, 2] with n=3: prefix [1, 2]; existing + # trigram (1, 2, 3) -> ban 3. + req = self.MockLlmRequest(tokens=[1, 2, 3, 4, 1, 2]) + logits = self._run([req], [3]) + assert self._banned_cols(logits[0]) == {3} + + def test_multiple_matches_all_banned(self): + # Prefix [2] occurs twice with different continuations -> ban both. + req = self.MockLlmRequest(tokens=[2, 3, 2, 5, 2]) + logits = self._run([req], [2]) + assert self._banned_cols(logits[0]) == {3, 5} + + def test_no_match_nothing_banned(self): + req = self.MockLlmRequest(tokens=[1, 2, 3, 4]) + logits = self._run([req], [3]) + assert self._banned_cols(logits[0]) == set() + + def test_sequence_shorter_than_ngram(self): + req = self.MockLlmRequest(tokens=[1, 2]) + logits = self._run([req], [3]) + assert self._banned_cols(logits[0]) == set() + + def test_unigram_bans_all_seen_tokens(self): + req = self.MockLlmRequest(tokens=[4, 7, 4]) + logits = self._run([req], [1]) + assert self._banned_cols(logits[0]) == {4, 7} + + @pytest.mark.parametrize("disabled_size", [None, 0]) + def test_disabled_request_untouched(self, disabled_size): + # Both None and 0 disable the restriction for that request. + active = self.MockLlmRequest(tokens=[1, 2, 3, 1, 2]) + disabled = self.MockLlmRequest(tokens=[1, 2, 3, 1, 2]) + logits = self._run([active, disabled], [2, disabled_size]) + assert self._banned_cols(logits[0]) == {3} + assert self._banned_cols(logits[1]) == set() + + def test_incremental_cache_updates_with_growing_context(self): + # The per-request index must extend as the token history grows. + # Reuse one sampler so the incremental cache carries across calls. + handler = SynchronousTokenBanHandler() + req = self.MockLlmRequest(tokens=[1, 2, 3]) + logits = self._run([req], [2], handler=handler) + assert self._banned_cols(logits[0]) == set() + req._tokens.extend([1, 2]) # now [1, 2, 3, 1, 2]: bigram (2, 3) -> ban 3 + logits = self._run([req], [2], handler=handler) + assert self._banned_cols(logits[0]) == {3} + + def test_cache_rebuilt_on_history_rollback(self): + # A shrunken history (speculative rollback) must invalidate the cache. + # Reuse one sampler so the stale cache would be hit if not rebuilt. + handler = SynchronousTokenBanHandler() + req = self.MockLlmRequest(tokens=[1, 2, 3, 1, 2]) + logits = self._run([req], [2], handler=handler) + assert self._banned_cols(logits[0]) == {3} + req._tokens[:] = [4, 5] + logits = self._run([req], [2], handler=handler) + assert self._banned_cols(logits[0]) == set() + + def test_multi_step_rows_all_banned(self): + # Speculative steps share the (host-approximated) banned set. + req = self.MockLlmRequest(tokens=[1, 2, 3, 1, 2]) + logits = self._run([req], [2], num_steps=[2]) + assert self._banned_cols(logits[0]) == {3} + assert self._banned_cols(logits[1]) == {3} + + def test_multi_beam_uses_per_beam_history(self): + # Beam histories diverge, so each beam's row is banned from its own + # n-grams; the shared per-request cache must not be used here. + req = self.MultiBeamMockLlmRequest( + beam_tokens=[ + [1, 2, 3, 1, 2], # beam 0: bigram (2, 3) -> ban 3 + [4, 5, 6, 4, 5], # beam 1: bigram (5, 6) -> ban 6 + ] + ) + # A beam-major / step-minor row layout with num_beams=2, num_steps=1 + # yields one row per beam. + logits = self._run([req], [2], num_beams=[2]) + assert self._banned_cols(logits[0]) == {3} + assert self._banned_cols(logits[1]) == {6} + + # --- Overlap-scheduler (stale-host) path ------------------------------- + # + # The host token list lags the device state by one token, read from + # new_tokens_cuda[0, seq_slot, 0] on the GPU. + + NUM_SLOTS = 4 + + def _new_tokens_cuda(self, slot_tokens): + buf = torch.full((1, self.NUM_SLOTS, 1), -1, dtype=torch.int32, device="cuda") + for slot, tok in slot_tokens.items(): + buf[0, slot, 0] = tok + return buf + + def _run_stale(self, requests, ngram_sizes, stale_by_one, slot_tokens, handler=None): + total_rows = len(requests) + logits = self._distinct_logits(total_rows, self.VOCAB) + before = logits.clone() + handler = handler or OverlappedTokenBanHandler() + num_steps = [1] * len(requests) + num_beams = [1] * len(requests) + bans = handler.generate_ban_list( + cast(list[LlmRequest], requests), + num_steps, + num_beams, + ngram_sizes, + stale_by_one=stale_by_one, + ) + handler.apply_ban_list(logits, bans, new_tokens_cuda=self._new_tokens_cuda(slot_tokens)) + self._assert_only_banned_changed(logits, before) + return logits + + def test_stale_bigram_device_hit(self): + # True sequence [1, 2, 3, 1] + device 2 == [1, 2, 3, 1, 2], n=2: + # window (2, 3) matches the prefix [d=2] -> ban 3. + req = self.MockLlmRequest(tokens=[1, 2, 3, 1], seq_slot=0) + logits = self._run_stale([req], [2], [True], {0: 2}) + assert self._banned_cols(logits[0]) == {3} + + def test_stale_bigram_device_miss(self): + # Device token 5 never occurred before -> nothing banned. + req = self.MockLlmRequest(tokens=[1, 2, 3, 1], seq_slot=0) + logits = self._run_stale([req], [2], [True], {0: 5}) + assert self._banned_cols(logits[0]) == set() + + def test_stale_trigram_host_and_device_hit(self): + # True sequence [1, 2, 3, 4, 1] + device 2, n=3: prefix [1, 2]; + # window (1, 2, 3) -> ban 3. + req = self.MockLlmRequest(tokens=[1, 2, 3, 4, 1], seq_slot=1) + logits = self._run_stale([req], [3], [True], {1: 2}) + assert self._banned_cols(logits[0]) == {3} + + def test_stale_trigram_host_miss(self): + # Prefix head is [4]; no earlier window starts with 4 followed by the + # device token -> nothing banned. + req = self.MockLlmRequest(tokens=[1, 2, 3, 4], seq_slot=1) + logits = self._run_stale([req], [3], [True], {1: 5}) + assert self._banned_cols(logits[0]) == set() + + def test_stale_window_ending_at_device_token(self): + # True sequence [7, 7] + device 7, n=2: window (7, 7) ends at the + # device token; the match forces d == context[-1] -> ban 7. + req = self.MockLlmRequest(tokens=[7, 7], seq_slot=0) + logits = self._run_stale([req], [2], [True], {0: 7}) + assert self._banned_cols(logits[0]) == {7} + + def test_stale_unigram_bans_host_and_device_tokens(self): + req = self.MockLlmRequest(tokens=[4, 7], seq_slot=2) + logits = self._run_stale([req], [1], [True], {2: 9}) + assert self._banned_cols(logits[0]) == {4, 7, 9} + + def test_stale_and_fresh_requests_mixed(self): + stale_req = self.MockLlmRequest(tokens=[1, 2, 3, 1], seq_slot=0) + fresh_req = self.MockLlmRequest(tokens=[1, 2, 3, 1, 2], seq_slot=1) + logits = self._run_stale([stale_req, fresh_req], [2, 2], [True, False], {0: 2}) + assert self._banned_cols(logits[0]) == {3} + assert self._banned_cols(logits[1]) == {3} + + def test_stale_no_nan_on_duplicate_bans(self): + # Duplicate bans on one cell must combine to -inf, not NaN: bigrams + # starting with 2 continue with 3 (twice) and with 2 (window at d). + req = self.MockLlmRequest(tokens=[2, 3, 2, 3, 2], seq_slot=0) + logits = self._run_stale([req], [2], [True], {0: 2}) + assert self._banned_cols(logits[0]) == {2, 3} + assert not torch.isnan(logits).any() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +class TestAddMinLengthBans: + """Unit tests for the min-length EOS-suppression path of TokenBanHandler. + + EOS (the request's original end_id) is banned on every step whose generated + length is still below py_min_length; it is an unconditional ban with no + stale-host variant. + """ + + VOCAB = 16 + END_ID = 5 + + class MockLlmRequest: + """Minimal stub exposing the attributes the min-length path reads.""" + + def __init__(self, num_tokens, *, min_length=None, prompt_len=0, end_id=5): + self._num_tokens = num_tokens # total tokens (prompt + generated) + self.py_min_length = [min_length] if min_length is not None else None + self.py_orig_prompt_len = prompt_len + self.py_original_end_id = end_id + self.py_end_id = end_id + + def get_num_tokens(self, beam_idx): + return self._num_tokens + + @staticmethod + def _distinct_logits(total_rows, vocab): + return (torch.arange(total_rows * vocab, device="cuda", dtype=torch.float) + 1.0).reshape( + total_rows, vocab + ) + + @staticmethod + def _banned_cols(logits_row): + return set(torch.nonzero(torch.isinf(logits_row)).flatten().tolist()) + + def _run(self, requests, num_steps=None, num_beams=None): + num_steps = num_steps or [1] * len(requests) + num_beams = num_beams or [1] * len(requests) + total_rows = sum(s * b for s, b in zip(num_steps, num_beams, strict=True)) + logits = self._distinct_logits(total_rows, self.VOCAB) + before = logits.clone() + ngram_sizes: list[int | None] = [None] * len(requests) + handler = SynchronousTokenBanHandler() + bans = handler.generate_ban_list( + cast(list[LlmRequest], requests), num_steps, num_beams, ngram_sizes + ) + handler.apply_ban_list(logits, bans) + # Only banned cells (-inf) changed; everything else keeps its value. + banned = torch.isinf(logits) + assert torch.equal(logits[~banned], before[~banned]) + return logits + + def test_below_min_length_bans_eos(self): + # 2 generated tokens (num_tokens 2, prompt 0), min_length 5 -> ban EOS. + req = self.MockLlmRequest(num_tokens=2, min_length=5) + logits = self._run([req]) + assert self._banned_cols(logits[0]) == {self.END_ID} + + def test_at_min_length_no_ban(self): + # 5 generated tokens, min_length 5 -> already satisfied, nothing banned. + req = self.MockLlmRequest(num_tokens=5, min_length=5) + logits = self._run([req]) + assert self._banned_cols(logits[0]) == set() + + def test_generated_length_excludes_prompt(self): + # num_tokens 7 with prompt_len 4 -> only 3 generated < min_length 5 -> ban. + req = self.MockLlmRequest(num_tokens=7, min_length=5, prompt_len=4) + logits = self._run([req]) + assert self._banned_cols(logits[0]) == {self.END_ID} + + def test_no_min_length_untouched(self): + active = self.MockLlmRequest(num_tokens=1, min_length=5) + disabled = self.MockLlmRequest(num_tokens=1, min_length=None) + logits = self._run([active, disabled]) + assert self._banned_cols(logits[0]) == {self.END_ID} + assert self._banned_cols(logits[1]) == set() + + def test_invalid_end_id_skipped(self): + # end_id <= -1 (e.g. ignore_eos) -> nothing to suppress. + req = self.MockLlmRequest(num_tokens=1, min_length=5, end_id=-1) + logits = self._run([req]) + assert self._banned_cols(logits[0]) == set() + + def test_multi_step_stops_at_min_length(self): + # 3 generated, min_length 5, 3 speculative steps: gen+step < 5 for + # steps 0,1 (3,4) but not step 2 (5) -> ban only rows 0 and 1. + req = self.MockLlmRequest(num_tokens=3, min_length=5) + logits = self._run([req], num_steps=[3]) + assert self._banned_cols(logits[0]) == {self.END_ID} + assert self._banned_cols(logits[1]) == {self.END_ID} + assert self._banned_cols(logits[2]) == set() diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index 892b93357fb0..ec89f4e7ca4c 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -2627,144 +2627,6 @@ def _uut(res=res): run_test_with_warmup(_uut_provider, max_sync_s=0.2) -@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") -class TestApplyBadWords: - """Unit tests for TorchSampler._apply_bad_words. - - Single-token words are banned unconditionally; multi-token words ban their - final token only when the token suffix (prompt + generated) matches the - word prefix. - """ - - VOCAB = 16 - - class MockLlmRequest: - """Minimal stub exposing the attributes _apply_bad_words reads.""" - - def __init__(self, tokens, *, bad_words=None, prompt_len=0, seq_slot=None): - # get_tokens(beam) returns the full token sequence (prompt + - # generated); py_orig_prompt_len marks where generation starts. - self.py_orig_prompt_len = prompt_len - self._tokens = list(tokens) - self.py_bad_words = bad_words - self.py_seq_slot = seq_slot - - def get_tokens(self, beam_idx): - return self._tokens - - def _run(self, requests, num_steps, num_beams): - total_rows = sum(s * b for s, b in zip(num_steps, num_beams)) - logits = torch.zeros(total_rows, self.VOCAB, device="cuda") - TorchSampler._apply_bad_words( - logits, cast(list[LlmRequest], requests), num_steps, num_beams - ) - return logits - - @staticmethod - def _banned_cols(logits_row): - return set(torch.nonzero(torch.isinf(logits_row)).flatten().tolist()) - - def test_single_token_unconditional(self): - req = self.MockLlmRequest(tokens=[3, 4], bad_words=[[7]]) - logits = self._run([req], num_steps=[1], num_beams=[1]) - assert self._banned_cols(logits[0]) == {7} - - def test_multi_token_prefix_hit(self): - # tokens end with [9]; word [9, 2] -> ban token 2. - req = self.MockLlmRequest(tokens=[1, 9], bad_words=[[9, 2]]) - logits = self._run([req], num_steps=[1], num_beams=[1]) - assert self._banned_cols(logits[0]) == {2} - - def test_multi_token_prefix_miss(self): - # tokens end with [3]; word [9, 2] prefix [9] does not match. - req = self.MockLlmRequest(tokens=[1, 3], bad_words=[[9, 2]]) - logits = self._run([req], num_steps=[1], num_beams=[1]) - assert self._banned_cols(logits[0]) == set() - - def test_multi_token_prefix_in_prompt(self): - # The prefix [9] lies in the prompt (nothing generated yet); the word - # [9, 2] must still ban token 2. - req = self.MockLlmRequest(tokens=[1, 9], bad_words=[[9, 2]], prompt_len=2) - logits = self._run([req], num_steps=[1], num_beams=[1]) - assert self._banned_cols(logits[0]) == {2} - - # --- Overlap-scheduler (stale-host) path ------------------------------- - # - # With the overlap scheduler the host token list lags the device state by - # one token; the newest token is read from new_tokens_cuda[0, seq_slot, 0] - # on the GPU. These tests drive _apply_bad_words with stale_by_one set. - - NUM_SLOTS = 4 - - def _new_tokens_cuda(self, slot_tokens): - buf = torch.full((1, self.NUM_SLOTS, 1), -1, dtype=torch.int32, device="cuda") - for slot, tok in slot_tokens.items(): - buf[0, slot, 0] = tok - return buf - - def _run_stale(self, requests, stale_by_one, slot_tokens): - total_rows = len(requests) - logits = torch.zeros(total_rows, self.VOCAB, device="cuda") - TorchSampler._apply_bad_words( - logits, - cast(list[LlmRequest], requests), - [1] * len(requests), - [1] * len(requests), - new_tokens_cuda=self._new_tokens_cuda(slot_tokens), - stale_by_one=stale_by_one, - ) - return logits - - def test_stale_two_token_device_hit(self): - # Host context [1]; device holds the pending token 9; word [9, 2] - # completes its prefix on the device side -> ban token 2. - req = self.MockLlmRequest(tokens=[1], bad_words=[[9, 2]], seq_slot=0) - logits = self._run_stale([req], [True], {0: 9}) - assert self._banned_cols(logits[0]) == {2} - - def test_stale_two_token_device_miss(self): - # Device token is 3, not the required prefix 9 -> nothing banned. - req = self.MockLlmRequest(tokens=[1], bad_words=[[9, 2]], seq_slot=0) - logits = self._run_stale([req], [True], {0: 3}) - assert self._banned_cols(logits[0]) == set() - - def test_stale_three_token_host_and_device_hit(self): - # Word [5, 9, 2]: host suffix must be [5], device token must be 9. - req = self.MockLlmRequest(tokens=[1, 5], bad_words=[[5, 9, 2]], seq_slot=1) - logits = self._run_stale([req], [True], {1: 9}) - assert self._banned_cols(logits[0]) == {2} - - def test_stale_three_token_host_miss(self): - # Host suffix [3] does not match the word prefix [5]; the device token - # matching is irrelevant -> nothing banned. - req = self.MockLlmRequest(tokens=[1, 3], bad_words=[[5, 9, 2]], seq_slot=1) - logits = self._run_stale([req], [True], {1: 9}) - assert self._banned_cols(logits[0]) == set() - - def test_stale_single_token_unconditional(self): - # Single-token words are banned regardless of the device token. - req = self.MockLlmRequest(tokens=[1], bad_words=[[7]], seq_slot=0) - logits = self._run_stale([req], [True], {0: 3}) - assert self._banned_cols(logits[0]) == {7} - - def test_stale_and_fresh_requests_mixed(self): - # Request 0 (stale) matches via the device token; request 1 (fresh) - # matches via the host context, as on the regular path. - stale_req = self.MockLlmRequest(tokens=[1], bad_words=[[9, 2]], seq_slot=0) - fresh_req = self.MockLlmRequest(tokens=[1, 9], bad_words=[[9, 4]], seq_slot=1) - logits = self._run_stale([stale_req, fresh_req], [True, False], {0: 9}) - assert self._banned_cols(logits[0]) == {2} - assert self._banned_cols(logits[1]) == {4} - - def test_stale_conditional_and_unconditional_same_cell(self): - # An unconditional single-token ban and a device-conditional ban on the - # same logit must combine to -inf (no NaN from -inf + -inf). - req = self.MockLlmRequest(tokens=[1], bad_words=[[2], [9, 2]], seq_slot=0) - logits = self._run_stale([req], [True], {0: 9}) - assert self._banned_cols(logits[0]) == {2} - assert not torch.isnan(logits).any() - - class TestTopPDecay: """Minimal functional guards for Top-P Decay in TorchSampler. diff --git a/tests/unittest/llmapi/test_llm_pytorch.py b/tests/unittest/llmapi/test_llm_pytorch.py index eaadc09df8af..bbfbc3a89885 100644 --- a/tests/unittest/llmapi/test_llm_pytorch.py +++ b/tests/unittest/llmapi/test_llm_pytorch.py @@ -1187,7 +1187,7 @@ def test_min_tokens(use_speculative: bool): def test_min_tokens_long_prompt(): """Check min_tokens is respected when prompt is longer than min_tokens. - Regression test for NVBug 5823135: _apply_min_length_penalty compared + Regression test for NVBug 5823135: the min-length EOS suppression compared total token count (prompt + generated) against the raw min_tokens value instead of comparing generated token count only. When prompt_len >= min_tokens the EOS suppression was never activated, allowing early