[TRTLLM-13233][feat] Support no_repeat_ngram_size in TorchSampler#16594
[TRTLLM-13233][feat] Support no_repeat_ngram_size in TorchSampler#16594zhaoyangwang-nvidia wants to merge 6 commits into
no_repeat_ngram_size in TorchSampler#16594Conversation
zhaoyangwang-nvidia
changed the title
zhaoyangwang-nvidia
changed the title
no_repeat_ngram_size in TorchSampler
zhaoyangwang-nvidia
changed the title
no_repeat_ngram_size in TorchSamplerno_repeat_ngram_size in TorchSampler
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThe Torch executor now propagates ChangesNo-repeat n-gram sampling
Estimated code review effort: 4 (Complex) | ~45 minutes Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant ExecutorRequest
participant LlmRequest
participant TorchSampler
participant Logits
ExecutorRequest->>LlmRequest: propagate no_repeat_ngram_size
LlmRequest->>TorchSampler: provide request context and ngram size
TorchSampler->>TorchSampler: collect bad-word and no-repeat n-gram bans
TorchSampler->>Logits: apply accumulated token bans
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unittest/_torch/sampler/test_torch_sampler.py (1)
2789-2797: 📐 Maintainability & Code Quality | 🔵 TrivialConsider adding a multi-beam coverage case.
All
_run/_run_staleinvocations usenum_beams == 1, so the multi-beam branch of_collect_no_repeat_ngram_bans.fresh_rules(thenum_beams[index] != 1path that scans each beam's context directly rather than the incremental index) is never exercised. Since beam search is a supported production path, a follow-up test that passesnum_beams > 1with divergent per-beam histories would close this gap.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/sampler/test_torch_sampler.py` around lines 2789 - 2797, Add a focused multi-beam test using the existing _run or _run_stale helpers, passing num_beams greater than one with distinct histories for each beam. Assert the no-repeat n-gram bans for every beam, exercising _collect_no_repeat_ngram_bans.fresh_rules through its num_beams[index] != 1 path while preserving existing single-beam coverage.Source: Path instructions
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/unittest/_torch/sampler/test_torch_sampler.py`:
- Around line 2789-2797: Add a focused multi-beam test using the existing _run
or _run_stale helpers, passing num_beams greater than one with distinct
histories for each beam. Assert the no-repeat n-gram bans for every beam,
exercising _collect_no_repeat_ngram_bans.fresh_rules through its
num_beams[index] != 1 path while preserving existing single-beam coverage.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 8beb9dee-254c-4001-aea5-a5e5f79b7bfe
📒 Files selected for processing (5)
docs/source/features/sampling.mdtensorrt_llm/_torch/pyexecutor/llm_request.pytensorrt_llm/_torch/pyexecutor/sampler/sampler.pytensorrt_llm/sampling_params.pytests/unittest/_torch/sampler/test_torch_sampler.py
fd73666 to
7bf7eb4
Compare
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 <zhaoyangw@nvidia.com>
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 <zhaoyangw@nvidia.com>
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 <zhaoyangw@nvidia.com>
7bf7eb4 to
940c68e
Compare
|
Hi @lori-ren all comments were addressed, please take a look. |
940c68e to
646e258
Compare
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tests/unittest/_torch/sampler/test_torch_sampler.py (2)
2954-3024: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winFresh-path assertions verified correct.
Manually traced the index construction against all bigram/trigram/unigram/incremental/rollback/multi-step assertions in this block — all match the expected
_extend_no_repeat_ngram_indexsemantics.One gap:
test_disabled_request_untouched(Line 2989-2994) only coversngram_size=None. TheSamplingParams.no_repeat_ngram_sizecontract states bothNoneand0disable the restriction — worth adding a0case alongsideNoneto lock in that branch, since0exercises a different code path (an explicit falsy check) thanNone.✅ Suggested additional case
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]) + disabled_none = self.MockLlmRequest(tokens=[1, 2, 3, 1, 2]) + disabled_zero = self.MockLlmRequest(tokens=[1, 2, 3, 1, 2]) + logits = self._run([active, disabled_none, disabled_zero], [2, None, 0]) assert self._banned_cols(logits[0]) == {3} - assert self._banned_cols(logits[1]) == set() + assert self._banned_cols(logits[1]) == set() + assert self._banned_cols(logits[2]) == set()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/sampler/test_torch_sampler.py` around lines 2954 - 3024, Add coverage to test_disabled_request_untouched for an explicitly configured no_repeat_ngram_size of 0, alongside the existing None case. Verify the request remains unchanged and no tokens are banned, while preserving the current enabled-request assertion.
3025-3105: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winStale-host assertions verified correct; multi-beam divergence untested.
Traced the stale/device-boundary math for all 8 stale-path tests (including the subtle "virtual self-transition" case in
test_stale_no_nan_on_duplicate_bans) — all check out.
_make_sampler(Line 2926-2937) fixesmax_beam_width=1, and no test here passesnum_beams > 1with per-beam-divergent histories. Per_extend_no_repeat_ngram_index's own docstring, the incremental cache is single-beam only "as beam histories diverge," implying a separate/fallback code path handlesbeam_width > 1. That path isn't exercised in this file, despite the PR description listing "multi-beam layouts" among the 18 covered scenarios.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/sampler/test_torch_sampler.py` around lines 3025 - 3105, Add coverage for divergent multi-beam histories in the stale-host sampler tests; current _make_sampler fixes max_beam_width=1 and never exercises the fallback path. Create a sampler/request setup with num_beams greater than one and distinct per-beam token histories, then invoke _run_stale or the relevant sampling flow and assert each beam receives the correct no-repeat-ngram bans. Ensure the test validates the path used when _extend_no_repeat_ngram_index cannot use its single-beam incremental cache.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/unittest/_torch/sampler/test_torch_sampler.py`:
- Around line 2954-3024: Add coverage to test_disabled_request_untouched for an
explicitly configured no_repeat_ngram_size of 0, alongside the existing None
case. Verify the request remains unchanged and no tokens are banned, while
preserving the current enabled-request assertion.
- Around line 3025-3105: Add coverage for divergent multi-beam histories in the
stale-host sampler tests; current _make_sampler fixes max_beam_width=1 and never
exercises the fallback path. Create a sampler/request setup with num_beams
greater than one and distinct per-beam token histories, then invoke _run_stale
or the relevant sampling flow and assert each beam receives the correct
no-repeat-ngram bans. Ensure the test validates the path used when
_extend_no_repeat_ngram_index cannot use its single-beam incremental cache.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: dcf99299-557a-45d6-b3fa-5a910e2cba56
📒 Files selected for processing (5)
docs/source/features/sampling.mdtensorrt_llm/_torch/pyexecutor/llm_request.pytensorrt_llm/_torch/pyexecutor/sampler/sampler.pytensorrt_llm/sampling_params.pytests/unittest/_torch/sampler/test_torch_sampler.py
🚧 Files skipped from review as they are similar to previous changes (3)
- tensorrt_llm/sampling_params.py
- docs/source/features/sampling.md
- tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
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.<method> references flagged in review. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
dcead0f to
58f6d9d
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/_torch/sampler/test_torch_sampler.py (1)
3002-3007: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winCover the
0disabled value.The contract requires both
Noneand0to disable the restriction, but this test only exercisesNone. Add a second disabled request usingngram_size=0and assert that its logits remain unchanged.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/sampler/test_torch_sampler.py` around lines 3002 - 3007, Extend test_disabled_request_untouched to include a request configured with ngram_size=0, run it alongside the existing active and None-disabled requests, and assert its banned columns are empty. Preserve the existing assertions for the active request and the None-disabled request.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/_torch/sampler/test_torch_sampler.py`:
- Around line 2952-2960: Extend the tests around _run and _apply_no_repeat_ngram
with a num_beams=[2] case using distinct token histories for each beam.
Configure the mock request so get_tokens(beam_idx) returns beam-specific
sequences, then assert each logits row is banned only according to its
corresponding beam history.
---
Outside diff comments:
In `@tests/unittest/_torch/sampler/test_torch_sampler.py`:
- Around line 3002-3007: Extend test_disabled_request_untouched to include a
request configured with ngram_size=0, run it alongside the existing active and
None-disabled requests, and assert its banned columns are empty. Preserve the
existing assertions for the active request and the None-disabled request.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: a326faf9-9f4a-43d5-a483-72fee52fadd5
📒 Files selected for processing (2)
tensorrt_llm/_torch/pyexecutor/sampler/sampler.pytests/unittest/_torch/sampler/test_torch_sampler.py
🚧 Files skipped from review as they are similar to previous changes (1)
- tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
646e258 to
dcead0f
Compare
|
|
||
|
|
||
| @dataclass(kw_only=True) | ||
| class _TokenBans: |
There was a problem hiding this comment.
The sampler.sampler submodule is already very big. Perhaps the token ban stuff could have its own submodule?
| if not active(index, r): | ||
| continue | ||
|
|
||
| if stale_by_one is not None and stale_by_one[index]: |
There was a problem hiding this comment.
Where is fresh_rules processed? Somewhere, we should probably check if overlap scheduling is enabled or not and then route accordingly in here. After encapsulating the token-ban-handling, this might be as concise as choosing between two polymorphic variants of the ban handler (one variant corresponds to the "fresh" and the other to the "stale" paths in the current feature branch).
| 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( |
There was a problem hiding this comment.
Could do row_idx, col_idx = torch.tensor([bans.rows, bans.cols], ...)...
There was a problem hiding this comment.
dev_row_idxanddev_slot_idxcan be combined in the same manner- same for
cond_row_idx,cond_col_idx,slot_idx
| neg_inf, | ||
| torch.zeros((), dtype=logits.dtype, device=device), | ||
| ) | ||
| logits.index_put_((cond_row_idx, cond_col_idx), penalty, accumulate=True) |
There was a problem hiding this comment.
Consider cond_row_idx_sel, cond_col_idx_sel = cond_rowcol_idx[:, prev_tokens == expected], logits.index_put_((cond_row_idx_sel, cond_col_idx_sel), neginf)
| new_tokens_cuda=new_tokens_cuda, | ||
| stale_by_one=stale_by_one, | ||
| ) | ||
| req_num_steps_list = sampling_requests_metadata.req_num_steps.tolist() |
There was a problem hiding this comment.
Only call .tolist() if needed (cf. following ifs)
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 <zhaoyangw@nvidia.com>
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 <zhaoyangw@nvidia.com>
3 participants
Summary
Dev Engineer Review
no_repeat_ngram_sizesupport toTorchSampler, aligning with C++banRepeatNgramand Hugging Face behavior:n-gram in the generated sequence including prompt tokens.n == 1, bans all previously seen tokens (unigram repeat prevention).Noneor0._apply_token_bans._TokenBansto accumulate ban targets across unconditional, overlap-conditioned, and device-token-column cases._NgramIndexCachefor efficient updates.neg_inftensor for overlap conditional penalties to avoid extra allocations._process_requeststo compute per-requestngram_sizesand apply both features via shared collection/flush when overlap scheduler tracking is enabled._NgramIndexCacheinWeakKeyDictionary[LlmRequest, _NgramIndexCache]so state is reclaimed when requests are freed (including rollback scenarios).executor_request_to_llm_requestnow readsexecutor_request.sampling_config.no_repeat_ngram_size, normalizes falsy/disabled values toNone, and setsllm_request.py_no_repeat_ngram_size.docs/source/features/sampling.mdto document prompt-inclusive n-gram banning and thatNone/0disables the restriction.SamplingParamsdocstring to explicitly describe the prompt-inclusive behavior and the default (None).QA Engineer Review
tests/unittest/_torch/sampler/test_torch_sampler.py:TestApplyBadWordsto use a new sampler helper for both fresh and stale-host paths.TestApplyNoRepeatNgramwith aMockLlmRequestand helpers to validate banned-logit columns and correct index/cache behavior.tests/integration/test_lists/,test-db/, orqa/: not verifiable from the provided context (no test-list changes available).Description
SamplingParams.no_repeat_ngram_sizewas previously only honored by the C++ executor path(
banRepeatNgramkernel viabanWordsLayer); on the PyTorch backend it was silently ignored.This PR implements it in
TorchSampler.Semantics
Same as the C++
banRepeatNgramkernel and HF'sNoRepeatNGramLogitsProcessor: withno_repeat_ngram_size = n, the lastn - 1tokens of the sequence (prompt included) form thecurrent prefix; the final token of every n-gram already present in the sequence whose first
n - 1tokens match that prefix is masked to-inf, so no n-gram is ever generated twice.n == 1bans every token already present. The restriction is per-beam.None/0disablesthe restriction (docs updated accordingly).
Implementation
Unified suffix-rule framework. Bad words and no-repeat ngram are both instances of one
abstraction — a rule
(prefix, col)that bans tokencolwhen the sequence ends withprefix(a bad word
[t0..tk-1]is the static rule(word[:-1], word[-1]); n-grams are rules derivedfrom the sequence itself). A shared driver
_collect_suffix_rule_bansnow owns:token
d(innew_tokens_cuda[0, seq_slot, 0]), so a rule's match is split into a host-sideprefix check plus a single device-side comparison of
dagainst a host-known expected token,resolved at flush time with no device-to-host sync. Even the n-gram window whose banned token
is
ditself reduces to a host-known column (the match condition forcesd == context[-1]).Single flush. Bans from all features accumulate into one
_TokenBansand flush with oneH2D transfer + one
index_put_per ban category (previously one set per feature).Host cost: incremental per-request index. A naive per-step window scan is
O(batch × history) on the host. Two measured bottlenecks are removed:
(n-1)-gram → next tokensindex (plus an(n-2)-grampair index for the stale path) iscached on the request and extended only with newly covered windows: steady-state cost is
O(new tokens) + an O(1) lookup. Rebuilt automatically on history shrinkage (speculative
rollback) or ngram-size change; multi-beam requests keep the direct scan (diverging beam
histories cannot share a per-request cache).
get_tokens()copies the whole history through the bindings (~40 µs at 2048 tokens, measured).The index keeps a host-side token mirror grown via
get_num_tokens/get_last_tokensscalaraccessors, and the driver hands rule generators a memoized context thunk fetched only when a
rule actually needs the full history — bad-words single-token rules and the single-beam n-gram
path never pay the copy.
Performance
Isolated sampler benchmark (H200,
bs=32, vocab 152k, 2048-token context per request,n=3,p50 sampler-beyond-forward wall time, 1024 iters × 3 runs):
Limitations
be reconstructed by the single-token device comparison, so bans may be enforced one step late
(existing
warning_onceextended to cover both features).(same approximation as bad words).
Test Coverage
tests/unittest/_torch/sampler/test_torch_sampler.py::TestApplyNoRepeatNgram(18 cases):bigram/trigram hit & miss, multiple matches,
n == 1, prompt-internal n-grams, sequencesshorter than
n, disabled requests, multi-step (speculation) and multi-beam row layouts,the full stale-host overlap matrix (device hit/miss, window ending at the device token,
mixed stale+fresh batches, duplicate bans producing no NaN), incremental cache growth,
and cache rebuild on history rollback. All pass on H200.
cross-checked against a brute-force reference implementation (fresh path) and against
"fresh path on the completed sequence" (stale path), plus 300 randomized bad-words
equivalence cases guarding the shared-driver refactor.
PR Checklist
[JIRA/NVBUG/None][type] descriptionformatdocs/source/features/sampling.md,SamplingParamsdocstring)PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.