From a44d6035b42a9047139c5f0fa325a219d249506d Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Sun, 19 Jul 2026 20:00:46 -0700 Subject: [PATCH 1/8] Support no_repeat_ngram_size in TorchSampler Bans tokens that would recreate an n-gram already present in the sequence (prompt included), matching the C++ banRepeatNgram kernel and HF NoRepeatNGramLogitsProcessor semantics. Both bad words and no-repeat ngram reduce to suffix rules (prefix -> banned token); a shared driver owns packed-row bookkeeping, suffix matching, and the overlap-scheduler stale-host split (host prefix check + device token comparison). Bans from all features accumulate into one _TokenBans and flush with a single H2D transfer + index_put_ per category. Signed-off-by: ZhaoyangWang --- docs/source/features/sampling.md | 4 + tensorrt_llm/_torch/pyexecutor/llm_request.py | 7 + .../_torch/pyexecutor/sampler/sampler.py | 449 ++++++++++++++---- tensorrt_llm/sampling_params.py | 2 +- .../_torch/sampler/test_torch_sampler.py | 167 +++++++ 5 files changed, 537 insertions(+), 92 deletions(-) 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/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..7b9bf6ce8ec7 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -19,7 +19,7 @@ from collections.abc import Iterable, Iterator from concurrent import futures from contextlib import AbstractContextManager, contextmanager, nullcontext -from dataclasses import dataclass +from dataclasses import dataclass, field from itertools import repeat from typing import ( TYPE_CHECKING, @@ -1345,6 +1345,29 @@ def _record_sampler_event( ) +@dataclass(kw_only=True) +class _TokenBans: + """Accumulated (row, col) logit bans, flushed by ``_apply_token_bans``. + + Collectors for token-ban features (bad words, no-repeat ngram) append here + so that all bans reach the device with a single host-to-device transfer + per index list and a single ``index_put_`` per category. + """ + + # Unconditional bans: logits[row, col] = -inf. + rows: list[int] = field(default_factory=list) + cols: list[int] = field(default_factory=list) + # Overlap path: bans applied only where the device-side previous token + # (new_tokens_cuda[0, slot, 0]) equals the expected host token. + 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) + # Overlap path: the device-side previous token itself is the banned column. + dev_rows: list[int] = field(default_factory=list) + dev_slots: list[int] = field(default_factory=list) + + class TorchSampler(Sampler[SampleStateTorch], AsyncWorkerMixin): DEFAULT_MAX_STOP_WORD_LENGTH = 20 DEFAULT_MAX_STOP_WORDS = 10 @@ -4810,132 +4833,356 @@ def _apply_bad_words( 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. + """Collect and immediately flush bad-words bans; see _collect_bad_words_bans.""" + bans = _TokenBans() + TorchSampler._collect_bad_words_bans( + bans, requests, num_steps, num_beams, stale_by_one=stale_by_one + ) + TorchSampler._apply_token_bans(logits, bans, new_tokens_cuda=new_tokens_cuda) + + @staticmethod + def _collect_bad_words_bans( + bans: _TokenBans, + requests: list[LlmRequest], + num_steps: list[int], + num_beams: list[int], + *, + stale_by_one: list[bool] | None = None, + ) -> None: + """Collect "bad words" bans that mask 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 + The stale-host overlap path (see ``_collect_suffix_rule_bans``) needs + no bad-words-specific handling: a word's prefix and final token are + host-known by construction. + + 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, same layout as + ``_apply_min_length_penalty``). + requests: The requests, aligned with the packed logits rows. + num_steps: Number of steps per request. + num_beams: Number of beams per request. + 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. + """ + + def rules(index: int, r: LlmRequest, context: list[int]) -> Iterator[tuple[list[int], int]]: + # A bad word is the suffix rule "after word[:-1], ban word[-1]". + for word in getattr(r, "py_bad_words", None) or (): + if word: + yield word[:-1], word[-1] + + TorchSampler._collect_suffix_rule_bans( + bans, + requests, + num_steps, + num_beams, + active=lambda index, r: bool(getattr(r, "py_bad_words", None)), + fresh_rules=rules, + stale_rules=rules, + stale_by_one=stale_by_one, + ) + + @staticmethod + @torch.inference_mode() + def _apply_no_repeat_ngram( + logits: torch.Tensor, + requests: list[LlmRequest], + ngram_sizes: list[int | None], + num_steps: list[int], + num_beams: list[int], + *, + new_tokens_cuda: torch.Tensor | None = None, + stale_by_one: list[bool] | None = None, + ) -> None: + """Collect and immediately flush no-repeat-ngram bans; see _collect_no_repeat_ngram_bans.""" + bans = _TokenBans() + TorchSampler._collect_no_repeat_ngram_bans( + bans, requests, ngram_sizes, num_steps, num_beams, stale_by_one=stale_by_one + ) + TorchSampler._apply_token_bans(logits, bans, new_tokens_cuda=new_tokens_cuda) + + @staticmethod + def _collect_no_repeat_ngram_bans( + bans: _TokenBans, + requests: list[LlmRequest], + ngram_sizes: list[int | None], + num_steps: list[int], + num_beams: list[int], + *, + stale_by_one: list[bool] | None = None, + ) -> None: + """Collect bans for tokens that would repeat an existing n-gram. + + For a request with no-repeat-ngram size ``n``, the last ``n - 1`` tokens + of the sequence (prompt + generated) form the current prefix. Every + n-gram already present in the sequence whose first ``n - 1`` tokens + equal that prefix has its final token masked to -inf, so no n-gram is + ever generated twice (same semantics as the C++ ``banRepeatNgram`` + kernel and HF's ``NoRepeatNGramLogitsProcessor``). ``n == 1`` bans every + token already present in the sequence. The restriction is per-beam. + + On the stale-host overlap path (see ``_collect_suffix_rule_bans``) the + rules are enumerated over the true sequence ``context + [d]``: a window + whose banned token is ``d`` itself reduces to the host column + ``context[-1]`` (forced by the match condition), and ``n == 1`` + additionally bans ``d`` via a gathered column index (``stale_extra``). + + 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, same layout as + ``_collect_bad_words_bans``). + requests: The requests, aligned with the packed logits rows. + ngram_sizes: Effective no-repeat-ngram size per request; None + disables the restriction for that request. + num_steps: Number of steps per request. + num_beams: Number of beams per request. + 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. + """ + + def fresh_rules( + index: int, r: LlmRequest, context: list[int] + ) -> Iterator[tuple[list[int], int]]: + # Every existing n-gram is the suffix rule "after its first n - 1 + # tokens, ban its last token". + n = ngram_sizes[index] + length = len(context) + if n is None or length < n: + return + if n == 1: + for col in set(context): + yield [], col + return + for i in range(length - n + 1): + yield context[i : i + n - 1], context[i + n - 1] + + def stale_rules( + index: int, r: LlmRequest, context: list[int] + ) -> Iterator[tuple[list[int], int]]: + # Rules over the true sequence context + [d]. Windows whose banned + # token is d itself reduce to a host column: the match condition + # forces d == context[-1] there. + n = ngram_sizes[index] + length = len(context) + if n is None: + return + if n == 1: + # Host tokens; d itself is banned via stale_extra below. + for col in set(context): + yield [], col + return + for i in range(length - n + 2): + col = context[i + n - 1] if i + n - 1 < length else context[-1] + yield context[i : i + n - 1], col + + 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) + + TorchSampler._collect_suffix_rule_bans( + bans, + requests, + num_steps, + num_beams, + active=lambda index, r: ngram_sizes[index] is not None, + fresh_rules=fresh_rules, + stale_rules=stale_rules, + stale_extra=stale_extra, + stale_by_one=stale_by_one, + ) + + @staticmethod + def _collect_suffix_rule_bans( + bans: _TokenBans, + requests: list[LlmRequest], + num_steps: list[int], + num_beams: list[int], + *, + active: Callable[[int, LlmRequest], bool], + fresh_rules: Callable[[int, LlmRequest, list[int]], Iterable[tuple[list[int], int]]], + stale_rules: Callable[[int, LlmRequest, list[int]], Iterable[tuple[list[int], int]]], + 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: - 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. + bans: Accumulator the collected bans are appended to. 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. + active: ``(index, request) -> bool``; requests where this is False + are skipped entirely. + fresh_rules: ``(index, request, context) -> iterable`` of + ``(prefix, col)`` rules matched against the given beam context. + stale_rules: Like ``fresh_rules`` but matched against the true + sequence ``context + [d]``, where ``d`` is the device-side + previous token; every yielded prefix and col 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. """ - 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: + 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 bad-words path only supports a single step and beam" + "stale-host token-ban 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]]. + # Host context is missing the previous step's token d; the true + # sequence is context + [new_tokens_cuda[0, seq_slot, 0]]. The + # match of a rule prefix against that sequence splits into a + # host check of prefix[:-1] and a device comparison of d + # against prefix[-1], resolved at flush time. context = r.get_tokens(0) - for word in bad_words: - k = len(word) + seen_uncond: set[int] = set() + seen_cond: set[tuple[int, int]] = set() + for prefix, col in stale_rules(index, r, context): + k = len(prefix) if k == 0: + if col not in seen_uncond: + seen_uncond.add(col) + bans.rows.append(request_offset) + bans.cols.append(col) continue - if k == 1: - # Single-token word: banned unconditionally. - rows.append(request_offset) - cols.append(word[0]) + # True sequence length is len(context) + 1; need >= k. + if len(context) < k - 1: continue - # True sequence length is len(context) + 1; need >= k - 1. - if len(context) < k - 2: + if k > 1 and context[-(k - 1) :] != prefix[:-1]: continue - # Host part: all prefix tokens except the newest one. - if k > 2 and context[-(k - 2) :] != word[: k - 2]: + key = (prefix[-1], col) + if key in seen_cond: 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]) + 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]): # Full token sequence for this beam (prompt + generated), so a - # bad-word prefix ending inside the prompt is also matched. + # rule 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: + seen_cols: set[int] = set() + for prefix, col in fresh_rules(index, r, context): + if col in seen_cols: 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: + k = len(prefix) + if k > 0 and (len(context) < k or context[-k:] != prefix): continue - # Apply to every step row of this beam. + 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]): - rows.append(request_offset + num_steps[index] * beam_idx + step) - cols.append(col) + bans.rows.append(request_offset + num_steps[index] * beam_idx + step) + bans.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( + @staticmethod + @torch.inference_mode() + def _apply_token_bans( + logits: torch.Tensor, + bans: _TokenBans, + *, + new_tokens_cuda: torch.Tensor | None = None, + ) -> None: + """Flush accumulated token bans to the logits in-place. + + Bans collected from all token-ban features (bad words, no-repeat + ngram) are transferred to the device together: one host-to-device copy + per index list and one ``index_put_`` per ban category, instead of one + set per feature. + + Args: + logits: Flattened ``[total_rows, vocab]`` logits. Modified in-place. + bans: The accumulated bans. + new_tokens_cuda: Device buffer holding the previous step's sampled + tokens, shape ``[max_tokens, max_num_sequences, max_beam_width]``. + Required when ``bans`` contains conditional or device-column + entries (overlap scheduler stale-host path). + """ + neg_inf = torch.full((), float("-inf"), dtype=logits.dtype, device=logits.device) + + if bans.rows: + row_idx = torch.tensor(bans.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( + col_idx = torch.tensor(bans.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: + if bans.dev_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 - ) + 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 + device = logits.device + 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( - cond_expected, dtype=new_tokens_cuda.dtype, pin_memory=prefer_pinned() + 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] @@ -4944,7 +5191,7 @@ def _apply_bad_words( # would force a device-to-host sync). penalty = torch.where( prev_tokens == expected, - torch.full((), float("-inf"), dtype=logits.dtype, device=device), + neg_inf, torch.zeros((), dtype=logits.dtype, device=device), ) logits.index_put_((cond_row_idx, cond_col_idx), penalty, accumulate=True) @@ -5261,7 +5508,12 @@ def _process_requests( sampling_requests_metadata.req_num_beams, ) - if any(getattr(r, "py_bad_words", 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_bad_words or has_no_repeat_ngram: stale_by_one: list[bool] | None = None if self._track_pending_steps and not self._is_draft_batch(sampling_requests): pending = [ @@ -5274,23 +5526,38 @@ def _process_requests( 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. + # the single-token device-side check; bans that depend + # on the token history 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.", + "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", ) - self._apply_bad_words( - logits_cuda, - sampling_requests, - sampling_requests_metadata.req_num_steps.tolist(), - sampling_requests_metadata.req_num_beams.tolist(), - new_tokens_cuda=new_tokens_cuda, - stale_by_one=stale_by_one, - ) + req_num_steps_list = sampling_requests_metadata.req_num_steps.tolist() + req_num_beams_list = sampling_requests_metadata.req_num_beams.tolist() + # Collect bans from all features into one accumulator so they are + # flushed with a single H2D transfer + index_put_ per category. + bans = _TokenBans() + if has_bad_words: + self._collect_bad_words_bans( + bans, + sampling_requests, + req_num_steps_list, + req_num_beams_list, + stale_by_one=stale_by_one, + ) + if has_no_repeat_ngram: + self._collect_no_repeat_ngram_bans( + bans, + sampling_requests, + ngram_sizes, + req_num_steps_list, + req_num_beams_list, + stale_by_one=stale_by_one, + ) + self._apply_token_bans(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/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_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index 892b93357fb0..40e19c1520e8 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -2893,3 +2893,170 @@ def test_reject_speculative_draft_tokens(self): ) # Same request without decay is accepted. sampler.validate_request(self._mock_request(SamplingParams(top_p=0.9), draft_tokens=[1, 2])) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +class TestApplyNoRepeatNgram: + """Unit tests for TorchSampler._apply_no_repeat_ngram. + + 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 _apply_no_repeat_ngram reads.""" + + def __init__(self, tokens, *, seq_slot=None): + self._tokens = list(tokens) + self.py_seq_slot = seq_slot + + def get_tokens(self, beam_idx): + return self._tokens + + def _run(self, requests, ngram_sizes, 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)) + logits = torch.zeros(total_rows, self.VOCAB, device="cuda") + TorchSampler._apply_no_repeat_ngram( + logits, cast(list[LlmRequest], requests), ngram_sizes, num_steps, num_beams + ) + return logits + + @staticmethod + def _banned_cols(logits_row): + return set(torch.nonzero(torch.isinf(logits_row)).flatten().tolist()) + + 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} + + def test_disabled_request_untouched(self): + active = self.MockLlmRequest(tokens=[1, 2, 3, 1, 2]) + disabled = self.MockLlmRequest(tokens=[1, 2, 3, 1, 2]) + logits = self._run([active, disabled], [2, None]) + assert self._banned_cols(logits[0]) == {3} + assert self._banned_cols(logits[1]) == 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} + + # --- 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): + total_rows = len(requests) + logits = torch.zeros(total_rows, self.VOCAB, device="cuda") + TorchSampler._apply_no_repeat_ngram( + logits, + cast(list[LlmRequest], requests), + ngram_sizes, + [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_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): + # Multiple windows banning the same cell must combine to -inf. True + # sequence [2, 3, 2, 3, 2, 2]: bigrams starting with 2 continue with + # 3 (twice, deduplicated) and with 2 (the window ending 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() From a4b2d1fcb72e662aaf69b0c925d0a751bf378192 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Sun, 19 Jul 2026 20:30:28 -0700 Subject: [PATCH 2/8] Serve no-repeat-ngram bans from an incremental per-request index The per-step window scan over the full token history was O(batch x history) on the host (~15ms/step at bs=32 with 2048-token contexts on H200, vs a ~128us sampler baseline). Two host-cost sources are removed: - Cache an (n-1)-gram -> next-tokens index (plus an (n-2)-gram pair index for the stale overlap path) on the request and extend it only with newly covered windows: steady-state O(new tokens) + O(1) lookup. - Avoid the full-history get_tokens() copy through the bindings (~40us at 2048 tokens per request per step): the index keeps a host-side token mirror grown via get_num_tokens/get_last_tokens scalar calls, and the suffix-rule driver hands generators a memoized context thunk fetched only when a rule actually needs the history. Multi-beam requests keep the direct scan: diverging beam histories cannot share the per-request cache. The cache is rebuilt on history shrinkage (speculative rollback) or ngram-size change. Signed-off-by: ZhaoyangWang --- .../_torch/pyexecutor/sampler/sampler.py | 216 ++++++++++++++---- .../_torch/sampler/test_torch_sampler.py | 29 ++- 2 files changed, 199 insertions(+), 46 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 7b9bf6ce8ec7..c9a615a77f8f 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -1345,6 +1345,24 @@ def _record_sampler_event( ) +def _memoized_context(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 (row, col) logit bans, flushed by ``_apply_token_bans``. @@ -4856,10 +4874,6 @@ def _collect_bad_words_bans( ``k-1`` most recently generated tokens exactly match the prefix ``[t0, ..., t_{k-2}]``. - The stale-host overlap path (see ``_collect_suffix_rule_bans``) needs - no bad-words-specific handling: a word's prefix and final token are - host-known by construction. - Args: bans: Accumulator the collected bans are appended to; row indices refer to the flattened ``[total_rows, vocab]`` logits layout @@ -4874,20 +4888,48 @@ def _collect_bad_words_bans( These bans need ``new_tokens_cuda`` when flushed. """ - def rules(index: int, r: LlmRequest, context: list[int]) -> Iterator[tuple[list[int], int]]: - # A bad word is the suffix rule "after word[:-1], ban word[-1]". + def fresh_rules( + index: int, r: LlmRequest, get_context: Callable[[], list[int]] + ) -> Iterator[tuple[list[int], int]]: + # A bad word is the suffix rule "after word[:-1], ban word[-1]"; + # matching is left to the driver. for word in getattr(r, "py_bad_words", None) or (): if word: yield word[:-1], word[-1] + def stale_rules( + index: int, r: LlmRequest, get_context: Callable[[], list[int]] + ) -> Iterator[tuple[list[int], int]]: + # Stale contract: yield only host-matched rules. The first k - 2 + # prefix tokens are checked here against the host context; the + # final prefix token is compared against the device token d at + # flush time. The context (a full-history copy) is fetched only + # when a multi-token word needs it. + 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: + yield [], 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 + yield word[:-1], word[-1] + TorchSampler._collect_suffix_rule_bans( bans, requests, num_steps, num_beams, active=lambda index, r: bool(getattr(r, "py_bad_words", None)), - fresh_rules=rules, - stale_rules=rules, + fresh_rules=fresh_rules, + stale_rules=stale_rules, stale_by_one=stale_by_one, ) @@ -4910,6 +4952,51 @@ def _apply_no_repeat_ngram( ) TorchSampler._apply_token_bans(logits, bans, new_tokens_cuda=new_tokens_cuda) + @staticmethod + def _extend_no_repeat_ngram_index( + 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)``: ``tokens`` is a host-side + mirror of the request's token history; ``fresh_idx`` maps each + ``(n - 1)``-gram of that history to the set of tokens that followed + it; ``stale_idx`` (only populated for ``n > 2``) maps each + ``(n - 2)``-gram to the set of ``(next token, token after that)`` + pairs. The mirror is grown with the ``get_num_tokens`` / + ``get_last_tokens`` scalar accessors because ``get_tokens()`` copies + the whole history through the bindings (~40us at 2048 tokens); each + window is indexed exactly once, so the steady-state per-step cost is + O(new tokens) instead of O(history length). The cache is rebuilt when + the history shrinks (e.g. speculative-decoding rollback) or ``n`` + changes. Only valid for single-beam requests: beam histories diverge + and cannot share one per-request cache. + """ + num_tokens = request.get_num_tokens(0) + cache = getattr(request, "_py_ngram_index_cache", None) + if cache is not None and cache[0] == n and len(cache[1]) <= num_tokens: + _, tokens, indexed_upto, fresh_idx, stale_idx = cache + 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, ngram-size change, or shrunken history. + tokens, indexed_upto, fresh_idx, stale_idx = request.get_tokens(0), 0, {}, {} + end = max(len(tokens) - n + 1, 0) + for i in range(indexed_upto, end): + fresh_idx.setdefault(tuple(tokens[i : i + n - 1]), set()).add(tokens[i + n - 1]) + if n > 2: + stale_idx.setdefault(tuple(tokens[i : i + n - 2]), set()).add( + (tokens[i + n - 2], tokens[i + n - 1]) + ) + request._py_ngram_index_cache = (n, tokens, end, fresh_idx, stale_idx) # type: ignore[attr-defined] + return tokens, fresh_idx, stale_idx + @staticmethod def _collect_no_repeat_ngram_bans( bans: _TokenBans, @@ -4953,13 +5040,28 @@ def _collect_no_repeat_ngram_bans( """ def fresh_rules( - index: int, r: LlmRequest, context: list[int] + index: int, r: LlmRequest, get_context: Callable[[], list[int]] ) -> Iterator[tuple[list[int], int]]: # Every existing n-gram is the suffix rule "after its first n - 1 # tokens, ban its last token". n = ngram_sizes[index] + if n is None: + return + if num_beams[index] == 1: + # Single beam: O(1) lookup in the incrementally built index — + # no history copy; matching is already done, so emit + # unconditional rules. + tokens, fresh_idx, _ = TorchSampler._extend_no_repeat_ngram_index(r, n) + if len(tokens) < n: + return + key = tuple(tokens[-(n - 1) :]) if n > 1 else () + for col in fresh_idx.get(key, ()): + yield [], col + return + # Multi-beam: histories diverge, so scan this beam's context. + context = get_context() length = len(context) - if n is None or length < n: + if length < n: return if n == 1: for col in set(context): @@ -4969,23 +5071,41 @@ def fresh_rules( yield context[i : i + n - 1], context[i + n - 1] def stale_rules( - index: int, r: LlmRequest, context: list[int] + index: int, r: LlmRequest, get_context: Callable[[], list[int]] ) -> Iterator[tuple[list[int], int]]: - # Rules over the true sequence context + [d]. Windows whose banned - # token is d itself reduce to a host column: the match condition - # forces d == context[-1] there. + # Rules over the true sequence tokens + [d], already host-matched + # (stale contract): only the device comparison of d against + # prefix[-1] remains, so a single-token prefix suffices. The stale + # path always has num_beams == 1, so the index cache applies. n = ngram_sizes[index] - length = len(context) if n is None: return + tokens, fresh_idx, stale_idx = TorchSampler._extend_no_repeat_ngram_index(r, n) if n == 1: # Host tokens; d itself is banned via stale_extra below. - for col in set(context): + for col in fresh_idx.get((), ()): yield [], col return - for i in range(length - n + 2): - col = context[i + n - 1] if i + n - 1 < length else context[-1] - yield context[i : i + n - 1], col + # Windows fully inside the host history, matched on their first + # n - 2 tokens (the device token d supplies the final prefix + # position, compared at flush time). + if n == 2: + # The (n - 2)-gram head is empty: every indexed bigram is a + # candidate; reconstruct the (expected, col) pairs from + # fresh_idx instead of keeping a degenerate second index. + 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) :]), ()) + for expected, col in pairs: + yield [expected], col + # The window ending at d itself (banned token is d); its match + # condition forces d == tokens[-1], a host-known column. Its + # host-side part (first n - 2 window tokens vs the current head) + # must be checked here. + if len(tokens) >= n - 1 and (n == 2 or tokens[-(n - 1) : -1] == tokens[-(n - 2) :]): + yield [tokens[-1]], tokens[-1] def stale_extra(bans: _TokenBans, request_offset: int, index: int, r: LlmRequest) -> None: if ngram_sizes[index] == 1: @@ -5014,8 +5134,12 @@ def _collect_suffix_rule_bans( num_beams: list[int], *, active: Callable[[int, LlmRequest], bool], - fresh_rules: Callable[[int, LlmRequest, list[int]], Iterable[tuple[list[int], int]]], - stale_rules: Callable[[int, LlmRequest, list[int]], Iterable[tuple[list[int], int]]], + fresh_rules: Callable[ + [int, LlmRequest, Callable[[], list[int]]], Iterable[tuple[list[int], int]] + ], + stale_rules: Callable[ + [int, LlmRequest, Callable[[], list[int]]], Iterable[tuple[list[int], int]] + ], stale_extra: Callable[[_TokenBans, int, int, LlmRequest], None] | None = None, stale_by_one: list[bool] | None = None, ) -> None: @@ -5044,12 +5168,18 @@ def _collect_suffix_rule_bans( num_beams: Number of beams per request. active: ``(index, request) -> bool``; requests where this is False are skipped entirely. - fresh_rules: ``(index, request, context) -> iterable`` of - ``(prefix, col)`` rules matched against the given beam context. - stale_rules: Like ``fresh_rules`` but matched against the true - sequence ``context + [d]``, where ``d`` is the device-side - previous token; every yielded prefix and col must consist of - host-known tokens. + fresh_rules: ``(index, request, get_context) -> iterable`` of + ``(prefix, col)`` rules matched by the driver against the beam + context. ``get_context()`` returns the beam's token history, + fetched (and memoized) on first use only — a full-history copy + through the bindings — so generators that don't need it keep + the step cheap. + stale_rules: Like ``fresh_rules``, but the rules apply to the true + sequence ``context + [d]`` (``d`` = device-side previous + token) and must already be host-matched by the generator: the + driver only resolves the device comparison of ``d`` against + ``prefix[-1]`` (an empty prefix bans unconditionally). Every + yielded prefix and col 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). @@ -5073,25 +5203,17 @@ def _collect_suffix_rule_bans( assert r.py_seq_slot is not None # Host context is missing the previous step's token d; the true # sequence is context + [new_tokens_cuda[0, seq_slot, 0]]. The - # match of a rule prefix against that sequence splits into a - # host check of prefix[:-1] and a device comparison of d - # against prefix[-1], resolved at flush time. - context = r.get_tokens(0) + # generator yields host-matched rules; the remaining 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, context): - k = len(prefix) - if k == 0: + for prefix, col in stale_rules(index, r, _memoized_context(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 - # True sequence length is len(context) + 1; need >= k. - if len(context) < k - 1: - continue - if k > 1 and context[-(k - 1) :] != prefix[:-1]: - continue key = (prefix[-1], col) if key in seen_cond: continue @@ -5105,16 +5227,20 @@ def _collect_suffix_rule_bans( continue for beam_idx in range(num_beams[index]): - # Full token sequence for this beam (prompt + generated), so a - # rule prefix ending inside the prompt is also matched. - context = r.get_tokens(beam_idx) + # The context is the beam's full token sequence (prompt + + # generated), so a rule prefix ending inside the prompt is + # also matched; it is fetched at most once and only if a rule + # actually needs it. + get_context = _memoized_context(r, beam_idx) seen_cols: set[int] = set() - for prefix, col in fresh_rules(index, r, context): + for prefix, col in fresh_rules(index, r, get_context): if col in seen_cols: continue k = len(prefix) - if k > 0 and (len(context) < k or context[-k:] != prefix): - continue + 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 diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index 40e19c1520e8..4baf4b7547a3 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -2914,7 +2914,14 @@ def __init__(self, tokens, *, seq_slot=None): self.py_seq_slot = seq_slot def get_tokens(self, beam_idx): - return self._tokens + # 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] def _run(self, requests, ngram_sizes, num_steps=None, num_beams=None): num_steps = num_steps or [1] * len(requests) @@ -2972,6 +2979,26 @@ def test_disabled_request_untouched(self): assert self._banned_cols(logits[0]) == {3} assert self._banned_cols(logits[1]) == set() + def test_incremental_cache_updates_with_growing_context(self): + # Same request object across steps: the per-request n-gram index must + # extend incrementally as the token history grows. + req = self.MockLlmRequest(tokens=[1, 2, 3]) + logits = self._run([req], [2]) + 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]) + assert self._banned_cols(logits[0]) == {3} + + def test_cache_rebuilt_on_history_rollback(self): + # A shrunken history (e.g. speculative rollback) must invalidate the + # cached index instead of serving stale bans. + req = self.MockLlmRequest(tokens=[1, 2, 3, 1, 2]) + logits = self._run([req], [2]) + assert self._banned_cols(logits[0]) == {3} + req._tokens[:] = [4, 5] + logits = self._run([req], [2]) + 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]) From 4984424cfaeb1abc376dedd88683e9dceecab2b2 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Mon, 20 Jul 2026 00:53:00 -0700 Subject: [PATCH 3/8] Tighten no-repeat-ngram comments and docstrings Collector docstrings no longer repeat the shared driver's argument documentation (the row-layout description now lives once, on _collect_suffix_rule_bans), and inline comments that re-narrated the docstrings are condensed. No functional change (AST-identical modulo docstrings). Signed-off-by: ZhaoyangWang --- .../_torch/pyexecutor/sampler/sampler.py | 130 ++++++------------ .../_torch/sampler/test_torch_sampler.py | 11 +- 2 files changed, 48 insertions(+), 93 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index c9a615a77f8f..0a2dc78913df 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -4874,18 +4874,8 @@ def _collect_bad_words_bans( ``k-1`` most recently generated tokens exactly match the prefix ``[t0, ..., t_{k-2}]``. - 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, same layout as - ``_apply_min_length_penalty``). - requests: The requests, aligned with the packed logits rows. - num_steps: Number of steps per request. - num_beams: Number of beams per request. - 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. + All arguments are forwarded to ``_collect_suffix_rule_bans``, which + documents them and the ban row layout. """ def fresh_rules( @@ -4900,11 +4890,8 @@ def fresh_rules( def stale_rules( index: int, r: LlmRequest, get_context: Callable[[], list[int]] ) -> Iterator[tuple[list[int], int]]: - # Stale contract: yield only host-matched rules. The first k - 2 - # prefix tokens are checked here against the host context; the - # final prefix token is compared against the device token d at - # flush time. The context (a full-history copy) is fetched only - # when a multi-token word needs it. + # Stale contract: check the first k - 2 prefix tokens here; the + # last one is compared against the device token d at flush time. context: list[int] | None = None for word in getattr(r, "py_bad_words", None) or (): k = len(word) @@ -4960,19 +4947,15 @@ def _extend_no_repeat_ngram_index( ]: """Incrementally maintained n-gram index for a single-beam request. - Returns ``(tokens, fresh_idx, stale_idx)``: ``tokens`` is a host-side - mirror of the request's token history; ``fresh_idx`` maps each - ``(n - 1)``-gram of that history to the set of tokens that followed - it; ``stale_idx`` (only populated for ``n > 2``) maps each - ``(n - 2)``-gram to the set of ``(next token, token after that)`` - pairs. The mirror is grown with the ``get_num_tokens`` / - ``get_last_tokens`` scalar accessors because ``get_tokens()`` copies - the whole history through the bindings (~40us at 2048 tokens); each - window is indexed exactly once, so the steady-state per-step cost is - O(new tokens) instead of O(history length). The cache is rebuilt when - the history shrinks (e.g. speculative-decoding rollback) or ``n`` - changes. Only valid for single-beam requests: beam histories diverge - and cannot share one per-request cache. + 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) or ``n`` + changes; single-beam only, as beam histories diverge. """ num_tokens = request.get_num_tokens(0) cache = getattr(request, "_py_ngram_index_cache", None) @@ -5017,26 +5000,13 @@ def _collect_no_repeat_ngram_bans( kernel and HF's ``NoRepeatNGramLogitsProcessor``). ``n == 1`` bans every token already present in the sequence. The restriction is per-beam. - On the stale-host overlap path (see ``_collect_suffix_rule_bans``) the - rules are enumerated over the true sequence ``context + [d]``: a window - whose banned token is ``d`` itself reduces to the host column - ``context[-1]`` (forced by the match condition), and ``n == 1`` - additionally bans ``d`` via a gathered column index (``stale_extra``). + 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``). - 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, same layout as - ``_collect_bad_words_bans``). - requests: The requests, aligned with the packed logits rows. - ngram_sizes: Effective no-repeat-ngram size per request; None - disables the restriction for that request. - num_steps: Number of steps per request. - num_beams: Number of beams per request. - 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. + ``ngram_sizes`` holds the effective per-request size (None = disabled); + the remaining arguments are forwarded to ``_collect_suffix_rule_bans``, + which documents them and the ban row layout. """ def fresh_rules( @@ -5048,9 +5018,7 @@ def fresh_rules( if n is None: return if num_beams[index] == 1: - # Single beam: O(1) lookup in the incrementally built index — - # no history copy; matching is already done, so emit - # unconditional rules. + # Single beam: O(1) index lookup, rules come out pre-matched. tokens, fresh_idx, _ = TorchSampler._extend_no_repeat_ngram_index(r, n) if len(tokens) < n: return @@ -5073,10 +5041,8 @@ def fresh_rules( def stale_rules( index: int, r: LlmRequest, get_context: Callable[[], list[int]] ) -> Iterator[tuple[list[int], int]]: - # Rules over the true sequence tokens + [d], already host-matched - # (stale contract): only the device comparison of d against - # prefix[-1] remains, so a single-token prefix suffices. The stale - # path always has num_beams == 1, so the index cache applies. + # Stale contract: rules are already host-matched; only the device + # comparison of d against prefix[-1] remains. n = ngram_sizes[index] if n is None: return @@ -5086,13 +5052,9 @@ def stale_rules( for col in fresh_idx.get((), ()): yield [], col return - # Windows fully inside the host history, matched on their first - # n - 2 tokens (the device token d supplies the final prefix - # position, compared at flush time). if n == 2: - # The (n - 2)-gram head is empty: every indexed bigram is a - # candidate; reconstruct the (expected, col) pairs from - # fresh_idx instead of keeping a degenerate second index. + # 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 ) @@ -5100,10 +5062,8 @@ def stale_rules( pairs = stale_idx.get(tuple(tokens[-(n - 2) :]), ()) for expected, col in pairs: yield [expected], col - # The window ending at d itself (banned token is d); its match - # condition forces d == tokens[-1], a host-known column. Its - # host-side part (first n - 2 window tokens vs the current head) - # must be checked here. + # 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) :]): yield [tokens[-1]], tokens[-1] @@ -5162,24 +5122,26 @@ def _collect_suffix_rule_bans( ``num_beams == 1`` (no speculation, no beam search). Args: - bans: Accumulator the collected bans are appended to. + 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, same layout as + ``_apply_min_length_penalty``). 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 matched by the driver against the beam - context. ``get_context()`` returns the beam's token history, - fetched (and memoized) on first use only — a full-history copy - through the bindings — so generators that don't need it keep - the step cheap. - stale_rules: Like ``fresh_rules``, but the rules apply to the true - sequence ``context + [d]`` (``d`` = device-side previous - token) and must already be host-matched by the generator: the - driver only resolves the device comparison of ``d`` against - ``prefix[-1]`` (an empty prefix bans unconditionally). Every - yielded prefix and col must consist of host-known tokens. + ``(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). @@ -5201,10 +5163,8 @@ def _collect_suffix_rule_bans( "stale-host token-ban 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 d; the true - # sequence is context + [new_tokens_cuda[0, seq_slot, 0]]. The - # generator yields host-matched rules; the remaining device - # comparison of d against prefix[-1] is resolved at flush time. + # 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, _memoized_context(r, 0)): @@ -5227,10 +5187,8 @@ def _collect_suffix_rule_bans( continue for beam_idx in range(num_beams[index]): - # The context is the beam's full token sequence (prompt + - # generated), so a rule prefix ending inside the prompt is - # also matched; it is fetched at most once and only if a rule - # actually needs it. + # The context (prompt + generated) is fetched at most once, + # and only if a rule with a non-empty prefix needs it. get_context = _memoized_context(r, beam_idx) seen_cols: set[int] = set() for prefix, col in fresh_rules(index, r, get_context): diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index 4baf4b7547a3..00d1226bbb4a 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -2980,8 +2980,7 @@ def test_disabled_request_untouched(self): assert self._banned_cols(logits[1]) == set() def test_incremental_cache_updates_with_growing_context(self): - # Same request object across steps: the per-request n-gram index must - # extend incrementally as the token history grows. + # The per-request index must extend as the token history grows. req = self.MockLlmRequest(tokens=[1, 2, 3]) logits = self._run([req], [2]) assert self._banned_cols(logits[0]) == set() @@ -2990,8 +2989,7 @@ def test_incremental_cache_updates_with_growing_context(self): assert self._banned_cols(logits[0]) == {3} def test_cache_rebuilt_on_history_rollback(self): - # A shrunken history (e.g. speculative rollback) must invalidate the - # cached index instead of serving stale bans. + # A shrunken history (speculative rollback) must invalidate the cache. req = self.MockLlmRequest(tokens=[1, 2, 3, 1, 2]) logits = self._run([req], [2]) assert self._banned_cols(logits[0]) == {3} @@ -3080,9 +3078,8 @@ def test_stale_and_fresh_requests_mixed(self): assert self._banned_cols(logits[1]) == {3} def test_stale_no_nan_on_duplicate_bans(self): - # Multiple windows banning the same cell must combine to -inf. True - # sequence [2, 3, 2, 3, 2, 2]: bigrams starting with 2 continue with - # 3 (twice, deduplicated) and with 2 (the window ending at d). + # 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} From 58f6d9d20784b645403bd05bd9763fc2f0ebca4e Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 22 Jul 2026 23:59:14 -0700 Subject: [PATCH 4/8] Address review: type the n-gram index cache and move it off LlmRequest Replace the untyped 5-tuple ngram index cache with a typed _NgramIndexCache dataclass, and store caches in a TorchSampler-owned WeakKeyDictionary keyed on the request instead of injecting a _py_ngram_index_cache attribute into LlmRequest. Weak keys let entries be reclaimed automatically when a request is freed, so no manual cleanup is needed. The no-repeat-ngram and bad-words token-ban helpers become instance methods and call each other and the shared suffix-rule driver via self, dropping the hardcoded TorchSampler. references flagged in review. Signed-off-by: ZhaoyangWang --- .../_torch/pyexecutor/sampler/sampler.py | 65 ++++++++++++++----- .../_torch/sampler/test_torch_sampler.py | 52 ++++++++++++--- 2 files changed, 89 insertions(+), 28 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 0a2dc78913df..3c5c8cf98251 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -34,6 +34,7 @@ TypeVar, cast, ) +from weakref import WeakKeyDictionary import numpy as np import torch @@ -1386,6 +1387,24 @@ class _TokenBans: dev_slots: list[int] = field(default_factory=list) +@dataclass +class _NgramIndexCache: + """Incrementally maintained n-gram index for a single request (single beam). + + ``fresh_idx`` maps each ``(n - 1)``-gram to the tokens that followed it, and + ``stale_idx`` (populated only for ``n > 2``) maps each ``(n - 2)``-gram to its + ``(next, next-but-one)`` token pairs. ``tokens`` is a host-side mirror of the + request's token history and ``indexed_upto`` records how far into it the + indices have been built, so each window is indexed exactly once. + """ + + n: int + tokens: list[int] + indexed_upto: int + fresh_idx: dict[tuple[int, ...], set[int]] + stale_idx: dict[tuple[int, ...], set[tuple[int, int]]] + + class TorchSampler(Sampler[SampleStateTorch], AsyncWorkerMixin): DEFAULT_MAX_STOP_WORD_LENGTH = 20 DEFAULT_MAX_STOP_WORDS = 10 @@ -2503,6 +2522,14 @@ def __init__(self, args: Args): # buffer entries are never consumed. self._top_p_decay_slots: set[int] = set() + # Per-request incrementally maintained n-gram indices for the + # no-repeat-ngram feature. Keyed weakly on the request so entries are + # reclaimed automatically when the request is freed (no manual cleanup, + # and no attribute injected into LlmRequest). + self._ngram_index_caches: WeakKeyDictionary[LlmRequest, _NgramIndexCache] = ( + WeakKeyDictionary() + ) + # AutoDeploy build creates the sampler in inference mode, # which would disallow in-place mutating of new_tokens. # So, we temporarily exit inference mode. @@ -4840,9 +4867,9 @@ def _apply_min_length_penalty( ) logits.index_put_((row_idx, col_idx), neg_inf, accumulate=False) - @staticmethod @torch.inference_mode() def _apply_bad_words( + self, logits: torch.Tensor, requests: list[LlmRequest], num_steps: list[int], @@ -4853,13 +4880,13 @@ def _apply_bad_words( ) -> None: """Collect and immediately flush bad-words bans; see _collect_bad_words_bans.""" bans = _TokenBans() - TorchSampler._collect_bad_words_bans( + self._collect_bad_words_bans( bans, requests, num_steps, num_beams, stale_by_one=stale_by_one ) - TorchSampler._apply_token_bans(logits, bans, new_tokens_cuda=new_tokens_cuda) + self._apply_token_bans(logits, bans, new_tokens_cuda=new_tokens_cuda) - @staticmethod def _collect_bad_words_bans( + self, bans: _TokenBans, requests: list[LlmRequest], num_steps: list[int], @@ -4909,7 +4936,7 @@ def stale_rules( continue yield word[:-1], word[-1] - TorchSampler._collect_suffix_rule_bans( + self._collect_suffix_rule_bans( bans, requests, num_steps, @@ -4920,9 +4947,9 @@ def stale_rules( stale_by_one=stale_by_one, ) - @staticmethod @torch.inference_mode() def _apply_no_repeat_ngram( + self, logits: torch.Tensor, requests: list[LlmRequest], ngram_sizes: list[int | None], @@ -4934,14 +4961,13 @@ def _apply_no_repeat_ngram( ) -> None: """Collect and immediately flush no-repeat-ngram bans; see _collect_no_repeat_ngram_bans.""" bans = _TokenBans() - TorchSampler._collect_no_repeat_ngram_bans( + self._collect_no_repeat_ngram_bans( bans, requests, ngram_sizes, num_steps, num_beams, stale_by_one=stale_by_one ) - TorchSampler._apply_token_bans(logits, bans, new_tokens_cuda=new_tokens_cuda) + self._apply_token_bans(logits, bans, new_tokens_cuda=new_tokens_cuda) - @staticmethod def _extend_no_repeat_ngram_index( - request: LlmRequest, n: int + self, request: LlmRequest, n: int ) -> tuple[ list[int], dict[tuple[int, ...], set[int]], dict[tuple[int, ...], set[tuple[int, int]]] ]: @@ -4958,9 +4984,10 @@ def _extend_no_repeat_ngram_index( changes; single-beam only, as beam histories diverge. """ num_tokens = request.get_num_tokens(0) - cache = getattr(request, "_py_ngram_index_cache", None) - if cache is not None and cache[0] == n and len(cache[1]) <= num_tokens: - _, tokens, indexed_upto, fresh_idx, stale_idx = cache + cache = self._ngram_index_caches.get(request) + if cache is not None and cache.n == n and len(cache.tokens) <= num_tokens: + tokens, indexed_upto = cache.tokens, cache.indexed_upto + 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)) @@ -4977,11 +5004,13 @@ def _extend_no_repeat_ngram_index( stale_idx.setdefault(tuple(tokens[i : i + n - 2]), set()).add( (tokens[i + n - 2], tokens[i + n - 1]) ) - request._py_ngram_index_cache = (n, tokens, end, fresh_idx, stale_idx) # type: ignore[attr-defined] + self._ngram_index_caches[request] = _NgramIndexCache( + n=n, tokens=tokens, indexed_upto=end, fresh_idx=fresh_idx, stale_idx=stale_idx + ) return tokens, fresh_idx, stale_idx - @staticmethod def _collect_no_repeat_ngram_bans( + self, bans: _TokenBans, requests: list[LlmRequest], ngram_sizes: list[int | None], @@ -5019,7 +5048,7 @@ def fresh_rules( return if num_beams[index] == 1: # Single beam: O(1) index lookup, rules come out pre-matched. - tokens, fresh_idx, _ = TorchSampler._extend_no_repeat_ngram_index(r, n) + tokens, fresh_idx, _ = self._extend_no_repeat_ngram_index(r, n) if len(tokens) < n: return key = tuple(tokens[-(n - 1) :]) if n > 1 else () @@ -5046,7 +5075,7 @@ def stale_rules( n = ngram_sizes[index] if n is None: return - tokens, fresh_idx, stale_idx = TorchSampler._extend_no_repeat_ngram_index(r, n) + tokens, fresh_idx, stale_idx = self._extend_no_repeat_ngram_index(r, n) if n == 1: # Host tokens; d itself is banned via stale_extra below. for col in fresh_idx.get((), ()): @@ -5074,7 +5103,7 @@ def stale_extra(bans: _TokenBans, request_offset: int, index: int, r: LlmRequest bans.dev_rows.append(request_offset) bans.dev_slots.append(r.py_seq_slot) - TorchSampler._collect_suffix_rule_bans( + self._collect_suffix_rule_bans( bans, requests, num_steps, diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index 00d1226bbb4a..4b3960faee88 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -2652,10 +2652,23 @@ def __init__(self, tokens, *, bad_words=None, prompt_len=0, seq_slot=None): def get_tokens(self, beam_idx): return self._tokens + @staticmethod + def _make_sampler(): + return TorchSampler( + TorchSampler.Args( + max_seq_len=128, + max_draft_len=0, + max_num_sequences=8, + max_beam_width=1, + max_total_draft_tokens=0, + disable_overlap_scheduler=True, + ) + ) + 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( + self._make_sampler()._apply_bad_words( logits, cast(list[LlmRequest], requests), num_steps, num_beams ) return logits @@ -2705,7 +2718,7 @@ def _new_tokens_cuda(self, slot_tokens): 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( + self._make_sampler()._apply_bad_words( logits, cast(list[LlmRequest], requests), [1] * len(requests), @@ -2923,12 +2936,26 @@ def get_num_tokens(self, beam_idx): def get_last_tokens(self, beam_idx): return self._tokens[-1] - def _run(self, requests, ngram_sizes, num_steps=None, num_beams=None): + @staticmethod + def _make_sampler(): + return TorchSampler( + TorchSampler.Args( + max_seq_len=128, + max_draft_len=0, + max_num_sequences=8, + max_beam_width=1, + max_total_draft_tokens=0, + disable_overlap_scheduler=True, + ) + ) + + def _run(self, requests, ngram_sizes, num_steps=None, num_beams=None, sampler=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)) logits = torch.zeros(total_rows, self.VOCAB, device="cuda") - TorchSampler._apply_no_repeat_ngram( + sampler = sampler or self._make_sampler() + sampler._apply_no_repeat_ngram( logits, cast(list[LlmRequest], requests), ngram_sizes, num_steps, num_beams ) return logits @@ -2981,20 +3008,24 @@ def test_disabled_request_untouched(self): 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. + sampler = self._make_sampler() req = self.MockLlmRequest(tokens=[1, 2, 3]) - logits = self._run([req], [2]) + logits = self._run([req], [2], sampler=sampler) 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]) + logits = self._run([req], [2], sampler=sampler) 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. + sampler = self._make_sampler() req = self.MockLlmRequest(tokens=[1, 2, 3, 1, 2]) - logits = self._run([req], [2]) + logits = self._run([req], [2], sampler=sampler) assert self._banned_cols(logits[0]) == {3} req._tokens[:] = [4, 5] - logits = self._run([req], [2]) + logits = self._run([req], [2], sampler=sampler) assert self._banned_cols(logits[0]) == set() def test_multi_step_rows_all_banned(self): @@ -3017,10 +3048,11 @@ def _new_tokens_cuda(self, slot_tokens): buf[0, slot, 0] = tok return buf - def _run_stale(self, requests, ngram_sizes, stale_by_one, slot_tokens): + def _run_stale(self, requests, ngram_sizes, stale_by_one, slot_tokens, sampler=None): total_rows = len(requests) logits = torch.zeros(total_rows, self.VOCAB, device="cuda") - TorchSampler._apply_no_repeat_ngram( + sampler = sampler or self._make_sampler() + sampler._apply_no_repeat_ngram( logits, cast(list[LlmRequest], requests), ngram_sizes, From def1459235110e12ea81f902b790f8b18773fff5 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Thu, 23 Jul 2026 02:47:19 -0700 Subject: [PATCH 5/8] Address review: naming, field doc-strings, and Final annotations Rename _memoized_context to _memoizing_tokens_getter (it returns a token getter, not a context manager). Give _NgramIndexCache per-field doc-strings instead of a long class doc-string, and rename its indexed_upto field to indexed_up_to. Annotate the DEFAULT_MAX_STOP_WORD_* class constants as Final[int]. Drop the confusing "as beam histories diverge" clause from the n-gram index doc-string and explain what "flush" means on _TokenBans. Signed-off-by: ZhaoyangWang --- .../_torch/pyexecutor/sampler/sampler.py | 43 ++++++++++--------- 1 file changed, 23 insertions(+), 20 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 3c5c8cf98251..986658be7f1f 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, @@ -1346,7 +1347,7 @@ def _record_sampler_event( ) -def _memoized_context(request: LlmRequest, beam_idx: int) -> Callable[[], list[int]]: +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 @@ -1370,7 +1371,9 @@ class _TokenBans: Collectors for token-ban features (bad words, no-repeat ngram) append here so that all bans reach the device with a single host-to-device transfer - per index list and a single ``index_put_`` per category. + per index list and a single ``index_put_`` per category. "Flushing" refers + to that final step: moving the accumulated host-side bans to the device and + applying them to the logits in one batched operation per category. """ # Unconditional bans: logits[row, col] = -inf. @@ -1389,25 +1392,25 @@ class _TokenBans: @dataclass class _NgramIndexCache: - """Incrementally maintained n-gram index for a single request (single beam). - - ``fresh_idx`` maps each ``(n - 1)``-gram to the tokens that followed it, and - ``stale_idx`` (populated only for ``n > 2``) maps each ``(n - 2)``-gram to its - ``(next, next-but-one)`` token pairs. ``tokens`` is a host-side mirror of the - request's token history and ``indexed_upto`` records how far into it the - indices have been built, so each window is indexed exactly once. - """ + """Incrementally maintained n-gram index for a single request (single beam).""" n: int + """N-gram size this index was built for; a change triggers a rebuild.""" tokens: list[int] - indexed_upto: 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: dict[tuple[int, ...], set[int]] + """Maps each ``(n - 1)``-gram to the tokens that followed it.""" stale_idx: dict[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 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 @@ -4981,12 +4984,12 @@ def _extend_no_repeat_ngram_index( ``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) or ``n`` - changes; single-beam only, as beam histories diverge. + changes; single-beam only. """ num_tokens = request.get_num_tokens(0) cache = self._ngram_index_caches.get(request) if cache is not None and cache.n == n and len(cache.tokens) <= num_tokens: - tokens, indexed_upto = cache.tokens, cache.indexed_upto + 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. @@ -4996,16 +4999,16 @@ def _extend_no_repeat_ngram_index( tokens = request.get_tokens(0) else: # First use, ngram-size change, or shrunken history. - tokens, indexed_upto, fresh_idx, stale_idx = request.get_tokens(0), 0, {}, {} + tokens, indexed_up_to, fresh_idx, stale_idx = request.get_tokens(0), 0, {}, {} end = max(len(tokens) - n + 1, 0) - for i in range(indexed_upto, end): + for i in range(indexed_up_to, end): fresh_idx.setdefault(tuple(tokens[i : i + n - 1]), set()).add(tokens[i + n - 1]) if n > 2: stale_idx.setdefault(tuple(tokens[i : i + n - 2]), set()).add( (tokens[i + n - 2], tokens[i + n - 1]) ) self._ngram_index_caches[request] = _NgramIndexCache( - n=n, tokens=tokens, indexed_upto=end, fresh_idx=fresh_idx, stale_idx=stale_idx + n=n, tokens=tokens, indexed_up_to=end, fresh_idx=fresh_idx, stale_idx=stale_idx ) return tokens, fresh_idx, stale_idx @@ -5196,7 +5199,7 @@ def _collect_suffix_rule_bans( # 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, _memoized_context(r, 0)): + 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) @@ -5218,7 +5221,7 @@ def _collect_suffix_rule_bans( 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 = _memoized_context(r, beam_idx) + 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: From 2c0eebedcbcf17927276d7fe253b5ddfe5a6da1b Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Thu, 23 Jul 2026 18:52:54 -0700 Subject: [PATCH 6/8] Use defaultdict(set) for the n-gram index setdefault(key, set()) builds a throwaway empty set on every call, even for keys already present. Switch the incrementally maintained fresh_idx / stale_idx to defaultdict(set) so a set is created only on first insertion of a key. All reads go through .get()/.items(), so the "read creates the key" behaviour of defaultdict is not triggered. Signed-off-by: ZhaoyangWang --- tensorrt_llm/_torch/pyexecutor/sampler/sampler.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 986658be7f1f..89626521a214 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -1401,9 +1401,9 @@ class _NgramIndexCache: indexed_up_to: int """How far into ``tokens`` the indices have been built, so each window is indexed exactly once.""" - fresh_idx: dict[tuple[int, ...], set[int]] + fresh_idx: defaultdict[tuple[int, ...], set[int]] """Maps each ``(n - 1)``-gram to the tokens that followed it.""" - stale_idx: dict[tuple[int, ...], set[tuple[int, int]]] + 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.""" @@ -4999,14 +4999,15 @@ def _extend_no_repeat_ngram_index( tokens = request.get_tokens(0) else: # First use, ngram-size change, or shrunken history. - tokens, indexed_up_to, fresh_idx, stale_idx = request.get_tokens(0), 0, {}, {} + 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): - fresh_idx.setdefault(tuple(tokens[i : i + n - 1]), set()).add(tokens[i + n - 1]) + # 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.setdefault(tuple(tokens[i : i + n - 2]), set()).add( - (tokens[i + n - 2], tokens[i + n - 1]) - ) + stale_idx[tuple(tokens[i : i + n - 2])].add((tokens[i + n - 2], tokens[i + n - 1])) self._ngram_index_caches[request] = _NgramIndexCache( n=n, tokens=tokens, indexed_up_to=end, fresh_idx=fresh_idx, stale_idx=stale_idx ) From 4fa0e472347cbc2145b1b49eafba351064173172 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Thu, 23 Jul 2026 23:27:04 -0700 Subject: [PATCH 7/8] Address review: return lists from rule generators, drop dead cache field, strengthen tests - fresh_rules / stale_rules (bad words and ngram) return list[...] instead of being generators: the driver always consumes them fully, so the generator's suspend/resume overhead has no lazy-evaluation payoff. - Drop the cache.n == n guard and the now-unused _NgramIndexCache.n field: a request's ngram size is fixed for its lifetime and the cache is keyed on the request, so the cached index always matches n. - Tests: cover the multi-beam per-beam-history path, the ngram_size == 0 disabled case, and assert non-banned logits are left untouched by starting from distinct non-zero logits; add zip(strict=True). Signed-off-by: ZhaoyangWang --- .../_torch/pyexecutor/sampler/sampler.py | 80 +++++++++---------- .../_torch/sampler/test_torch_sampler.py | 70 ++++++++++++++-- 2 files changed, 102 insertions(+), 48 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 89626521a214..15bf6290cc8f 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -1392,10 +1392,12 @@ class _TokenBans: @dataclass class _NgramIndexCache: - """Incrementally maintained n-gram index for a single request (single beam).""" + """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. + """ - n: int - """N-gram size this index was built for; a change triggers a rebuild.""" tokens: list[int] """Host-side mirror of the request's token history.""" indexed_up_to: int @@ -4910,25 +4912,26 @@ def _collect_bad_words_bans( def fresh_rules( index: int, r: LlmRequest, get_context: Callable[[], list[int]] - ) -> Iterator[tuple[list[int], int]]: + ) -> list[tuple[list[int], int]]: # A bad word is the suffix rule "after word[:-1], ban word[-1]"; # matching is left to the driver. - for word in getattr(r, "py_bad_words", None) or (): - if word: - yield word[:-1], word[-1] + 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]] - ) -> Iterator[tuple[list[int], int]]: + ) -> list[tuple[list[int], int]]: # 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[tuple[list[int], int]] = [] 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: - yield [], word[0] + rules.append(([], word[0])) continue if context is None: context = get_context() @@ -4937,7 +4940,8 @@ def stale_rules( continue if k > 2 and context[-(k - 2) :] != word[: k - 2]: continue - yield word[:-1], word[-1] + rules.append((word[:-1], word[-1])) + return rules self._collect_suffix_rule_bans( bans, @@ -4983,12 +4987,14 @@ def _extend_no_repeat_ngram_index( 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) or ``n`` - changes; single-beam only. + Rebuilt when the history shrinks (speculative rollback); single-beam + only. """ num_tokens = request.get_num_tokens(0) cache = self._ngram_index_caches.get(request) - if cache is not None and cache.n == n and len(cache.tokens) <= num_tokens: + # 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: @@ -4998,7 +5004,7 @@ def _extend_no_repeat_ngram_index( # Several tokens landed at once (e.g. accepted draft tokens). tokens = request.get_tokens(0) else: - # First use, ngram-size change, or shrunken history. + # 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) @@ -5009,7 +5015,7 @@ def _extend_no_repeat_ngram_index( 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( - n=n, tokens=tokens, indexed_up_to=end, fresh_idx=fresh_idx, stale_idx=stale_idx + tokens=tokens, indexed_up_to=end, fresh_idx=fresh_idx, stale_idx=stale_idx ) return tokens, fresh_idx, stale_idx @@ -5037,54 +5043,48 @@ def _collect_no_repeat_ngram_bans( ``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 = disabled); - the remaining arguments are forwarded to ``_collect_suffix_rule_bans``, + ``ngram_sizes`` holds the effective per-request size (None or a + non-positive value disables the restriction); the remaining arguments + are forwarded to ``_collect_suffix_rule_bans``, which documents them and the ban row layout. """ def fresh_rules( index: int, r: LlmRequest, get_context: Callable[[], list[int]] - ) -> Iterator[tuple[list[int], int]]: + ) -> list[tuple[list[int], int]]: # Every existing n-gram is the suffix rule "after its first n - 1 # tokens, ban its last token". n = ngram_sizes[index] - if n is None: - return + 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_no_repeat_ngram_index(r, n) if len(tokens) < n: - return + return [] key = tuple(tokens[-(n - 1) :]) if n > 1 else () - for col in fresh_idx.get(key, ()): - yield [], col - return + 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 + return [] if n == 1: - for col in set(context): - yield [], col - return - for i in range(length - n + 1): - yield context[i : i + n - 1], context[i + 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]] - ) -> Iterator[tuple[list[int], int]]: + ) -> list[tuple[list[int], int]]: # Stale contract: rules are already host-matched; only the device # comparison of d against prefix[-1] remains. n = ngram_sizes[index] - if n is None: - return + if not n: # None or non-positive: restriction disabled + return [] tokens, fresh_idx, stale_idx = self._extend_no_repeat_ngram_index(r, n) if n == 1: # Host tokens; d itself is banned via stale_extra below. - for col in fresh_idx.get((), ()): - yield [], col - return + 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. @@ -5093,12 +5093,12 @@ def stale_rules( ) else: pairs = stale_idx.get(tuple(tokens[-(n - 2) :]), ()) - for expected, col in pairs: - yield [expected], col + rules: list[tuple[list[int], int]] = [([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) :]): - yield [tokens[-1]], tokens[-1] + 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: @@ -5112,7 +5112,7 @@ def stale_extra(bans: _TokenBans, request_offset: int, index: int, r: LlmRequest requests, num_steps, num_beams, - active=lambda index, r: ngram_sizes[index] is not None, + active=lambda index, r: bool(ngram_sizes[index]), fresh_rules=fresh_rules, stale_rules=stale_rules, stale_extra=stale_extra, diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index 4b3960faee88..3f9a4d9f6b42 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -2936,34 +2936,69 @@ def get_num_tokens(self, beam_idx): 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 _make_sampler(): + def _make_sampler(max_beam_width=1): return TorchSampler( TorchSampler.Args( max_seq_len=128, max_draft_len=0, max_num_sequences=8, - max_beam_width=1, + max_beam_width=max_beam_width, max_total_draft_tokens=0, disable_overlap_scheduler=True, ) ) + @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, sampler=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)) - logits = torch.zeros(total_rows, self.VOCAB, device="cuda") - sampler = sampler or self._make_sampler() + 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() + sampler = sampler or self._make_sampler(max_beam_width=max(num_beams)) sampler._apply_no_repeat_ngram( logits, cast(list[LlmRequest], requests), ngram_sizes, num_steps, num_beams ) + 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. @@ -2999,10 +3034,12 @@ def test_unigram_bans_all_seen_tokens(self): logits = self._run([req], [1]) assert self._banned_cols(logits[0]) == {4, 7} - def test_disabled_request_untouched(self): + @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, None]) + logits = self._run([active, disabled], [2, disabled_size]) assert self._banned_cols(logits[0]) == {3} assert self._banned_cols(logits[1]) == set() @@ -3035,6 +3072,21 @@ def test_multi_step_rows_all_banned(self): 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 @@ -3050,7 +3102,8 @@ def _new_tokens_cuda(self, slot_tokens): def _run_stale(self, requests, ngram_sizes, stale_by_one, slot_tokens, sampler=None): total_rows = len(requests) - logits = torch.zeros(total_rows, self.VOCAB, device="cuda") + logits = self._distinct_logits(total_rows, self.VOCAB) + before = logits.clone() sampler = sampler or self._make_sampler() sampler._apply_no_repeat_ngram( logits, @@ -3061,6 +3114,7 @@ def _run_stale(self, requests, ngram_sizes, stale_by_one, slot_tokens, sampler=N new_tokens_cuda=self._new_tokens_cuda(slot_tokens), stale_by_one=stale_by_one, ) + self._assert_only_banned_changed(logits, before) return logits def test_stale_bigram_device_hit(self): From dcd3c22ff9a193ae5d128ab0314e4cd807235012 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Fri, 24 Jul 2026 00:03:23 -0700 Subject: [PATCH 8/8] Address review: build unconditional-ban row/col indices in one H2D transfer Pack bans.rows and bans.cols into a single [2, N] tensor and split it into two index rows on the device, so the host-to-device copy is one transfer instead of two. Signed-off-by: ZhaoyangWang --- tensorrt_llm/_torch/pyexecutor/sampler/sampler.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 15bf6290cc8f..408a8f5eb95d 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -5266,13 +5266,12 @@ def _apply_token_bans( neg_inf = torch.full((), float("-inf"), dtype=logits.dtype, device=logits.device) if bans.rows: - row_idx = torch.tensor(bans.rows, dtype=torch.long, pin_memory=prefer_pinned()).to( - logits.device, non_blocking=True - ) - col_idx = torch.tensor(bans.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) + # Build rows and cols as one [2, N] tensor so the host->device copy + # is a single transfer; split into two index rows on the 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) if bans.dev_rows: assert new_tokens_cuda is not None