[TRTLLM-13235][feat] Support bad_words in TorchSampler#16052
[TRTLLM-13235][feat] Support bad_words in TorchSampler#16052zhaoyangwang-nvidia wants to merge 4 commits into
Conversation
zhaoyangwang-nvidia
force-pushed
the
pr-15542-torch-badwords
branch
5 times, most recently
from
03266f1 to
07a4945
Compare
zhaoyangwang-nvidia
force-pushed
the
pr-15542-torch-badwords
branch
from
07a4945 to
294ee9b
Compare
|
/bot run --disable-fail-fast |
📝 WalkthroughWalkthroughThis PR adds bad-words masking to TorchSampler. LlmRequest now stores raw bad-words token lists via a new py_bad_words attribute, and TorchSampler applies a new _apply_bad_words method to mask corresponding logits to -inf during sampling, with unit tests covering single- and multi-token cases. ChangesBad Words Masking in TorchSampler
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant ProcessRequests as _process_requests
participant MinLength as _apply_min_length_penalty
participant BadWords as _apply_bad_words
participant Logits as logits_cuda
ProcessRequests->>MinLength: apply penalty(logits_cuda)
ProcessRequests->>BadWords: apply_bad_words(logits_cuda, py_bad_words)
BadWords->>Logits: set banned vocab columns to -inf
ProcessRequests->>ProcessRequests: proceed to greedy/strategy sampling
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Warning Review ran into problems🔥 ProblemsGit: Failed to clone repository. Please run the 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)
2695-2710: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winCoverage gap: speculative multi-step, multi-beam, and multi-request offsets are untested.
All three cases use
num_steps=[1], num_beams=[1]. The most defect-prone path — multi-token bad words withnum_steps > 1(speculative decoding), where per-step draft context matters — is not exercised, nor are multi-beam or multi-request row-offset layouts. Adding anum_steps=[2]multi-token case would surface the per-step context issue flagged insampler.py, and a two-request case would validatecurrent_offsetaccounting.Recommend follow-up tests in this file covering: speculative multi-step multi-token hit/miss,
num_beams>1per-beam banning, and two-request offset correctness.As per path instructions ("Act as a QA engineer reviewing test changes and coverage for TensorRT-LLM ... suggest ... whether coverage is sufficient, insufficient, or needs follow-up").
🤖 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 2695 - 2710, Add follow-up coverage in the existing bad-words sampler tests to exercise the uncovered paths in the sampler logic. Extend the test methods around test_single_token_unconditional, test_multi_token_prefix_hit, and test_multi_token_prefix_miss with cases for num_steps=[2] to verify speculative multi-step multi-token hit/miss behavior, and add assertions for num_beams>1 and a two-request scenario to validate per-beam banning and current_offset row handling. Keep the focus on the sampler behavior in test_torch_sampler.py so the cases hit the multi-step context path referenced by sampler.py.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 2695-2710: Add follow-up coverage in the existing bad-words
sampler tests to exercise the uncovered paths in the sampler logic. Extend the
test methods around test_single_token_unconditional,
test_multi_token_prefix_hit, and test_multi_token_prefix_miss with cases for
num_steps=[2] to verify speculative multi-step multi-token hit/miss behavior,
and add assertions for num_beams>1 and a two-request scenario to validate
per-beam banning and current_offset row handling. Keep the focus on the sampler
behavior in test_torch_sampler.py so the cases hit the multi-step context path
referenced by sampler.py.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 2bd5389c-6c50-4312-b217-e5ae1c547c9d
📒 Files selected for processing (3)
tensorrt_llm/_torch/pyexecutor/llm_request.pytensorrt_llm/_torch/pyexecutor/sampler/sampler.pytests/unittest/_torch/sampler/test_torch_sampler.py
|
PR_Github #58425 [ run ] triggered by Bot. Commit: |
|
PR_Github #58425 [ run ] completed with state
|
zhaoyangwang-nvidia
force-pushed
the
pr-15542-torch-badwords
branch
from
294ee9b to
4dcbb4e
Compare
|
/bot run |
|
PR_Github #58582 [ run ] triggered by Bot. Commit: |
|
PR_Github #58582 [ run ] completed with state
|
|
/bot run |
|
PR_Github #58651 [ run ] triggered by Bot. Commit: |
|
PR_Github #58651 [ run ] completed with state
|
|
/bot run |
|
Hi @NVIDIA/trt-llm-runtime-devs please help to reivew this PR, only small feature, thanks~ |
|
PR_Github #58948 [ run ] triggered by Bot. Commit: |
|
PR_Github #58948 [ run ] completed with state
|
zhaoyangwang-nvidia
force-pushed
the
pr-15542-torch-badwords
branch
from
4dcbb4e to
e8fd253
Compare
AFAIK stop words is applied after sampling, whereas bad words needs to be applied before sampling to adjust the logits. One could try to reuse or adjust |
The PyTorch TorchSampler previously ignored bad_words / bad_token_ids; only the C++ decoding path (TRTLLMSampler) applied them. Add logits masking in the TorchSampler so single- and multi-token bad words are banned, matching the C++ ban_bad_words kernel semantics. - llm_request: mirror the executor request's bad_words as py_bad_words (list[list[int]]) for the pure-Python sampler path. - sampler: add TorchSampler._apply_bad_words, called alongside the existing min-length penalty. Single-token words are banned unconditionally; multi-token words ban their final token only when the generated suffix matches the word prefix. Speculative steps reconstruct per-step context from py_draft_tokens up-front, matching the kernel's per-step progression. - tests: add TestApplyBadWords covering single/multi-token hit and miss, speculative multi-step, early-return, per-beam, and multi-request offsets. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
With the overlap scheduler, sample_async for step i runs before update_requests for step i-1, so LlmRequest.get_tokens() lags the device state by one token. _apply_bad_words matched multi-token bad-word prefixes against this stale host context, which could both miss bans (prefix completed by the previous step's token) and apply spurious ones. Split the prefix match for overlap-affected requests: the first k-2 prefix tokens are matched on the host, and the newest token is compared on the GPU against new_tokens[0, seq_slot, 0], applying -inf via an additive shape-static index_put_ so no device-to-host sync is added. TorchSampler now tracks, per seq slot, how many sampled steps have not been folded back into the request (incremented in sample_async, decremented in update_requests, reset when a context request claims a slot; draft batches excluded). The device-side path triggers only for overlap mode with pending == 1, no speculation and no beam search; speculative/beam overlap keeps the previous behavior and warns once. The non-overlap path is unchanged. Add TestApplyBadWords coverage for the stale-host path: device-side hit/miss, host+device combined match, mixed stale/fresh batches, and overlapping unconditional/conditional bans on the same logit. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
zhaoyangwang-nvidia
force-pushed
the
pr-15542-torch-badwords
branch
2 times, most recently
from
c83ea01 to
e241caf
Compare
Apply ixlmar's post-merge review suggestions from NVIDIA#16365: - ops/flashinfer.py: without flashinfer installed, calling a sampling op now raises an ImportError with installation guidance instead of a bare NameError. The module intentionally stays importable (sampling_utils imports it unconditionally and non-flashinfer sampler paths must keep working); startup-time enforcement remains in TorchSampler.__init__. Document that contract in the module docstring. - sampling_utils.py: collapse the "from x import y as y" re-export lines into plain grouped imports plus an explicit __all__ (mypy runs strict, so implicit re-export is disabled). - Drop the "greedy" alias of greedy_search_sampling_batch; use the descriptive name directly in drafting_loops.py and speculative/interface.py. - sanitize_top_k: replace torch.where(...).clamp(...) with clamp(...).masked_fill_(...), avoiding the full_like intermediate. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
|
Hi @ixlmar This PR addresses your post-merge review comments on #16365 (commit aded2c79):
|
The pending-step accounting classifies a batch as draft/target from its first request, assuming batches are homogeneous. That invariant is structural (ModelDrafter issues all-draft batches on the shared sampler; PyExecutor's are all-target) but was only documented. Add _is_draft_batch(), which asserts homogeneity, and use it at all three call sites so a mixed batch fails loudly instead of silently corrupting the counters. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #59693 [ run ] triggered by Bot. Commit: |
|
PR_Github #59693 [ run ] completed with state
|
6 participants
Summary by CodeRabbit
Description
TorchSamplerpreviously ignoredbad_words/bad_token_ids(only the C++TRTLLMSamplerapplied them). This PR adds enforcement, including correct behavior under the overlap scheduler.Bad-words masking in TorchSampler
llm_request.py: mirror the executor request'sbad_wordsaspy_bad_words.sampler.py: addTorchSampler._apply_bad_words— single-token words are banned unconditionally; a multi-token word bans its final token only when the generated suffix matches the word prefix (matches C++ban_bad_wordssemantics). Speculative steps reconstruct per-step context frompy_draft_tokens.Overlap scheduler support
Under overlap,
sample_async(i)runs beforeupdate_requests(i-1), so the host token list lags the device by one token — multi-token prefix matching could miss or misapply bans. Fix, with no added device-to-host sync:k-2prefix tokens on the host, the newest token compared on-GPU againststore.new_tokens[0, seq_slot, 0]; conditional ban applied via additive shape-staticindex_put_.Post-merge review follow-ups for #16365 (@ixlmar)
ops/flashinfer.py: missing flashinfer now raises an actionableImportErroron op use (module stays importable; startup enforcement stays inTorchSampler.__init__).sampling_utils.py: replaceimport y as yre-exports with plain imports +__all__; drop thegreedyalias (callers usegreedy_search_sampling_batch);sanitize_top_kusesclamp(...).masked_fill_(...).Follow-up (out of scope)
Batched device-side bad-words matching via the stop-words machinery (
past_tokens_cuda) would remove the host loop and cover spec/beam under overlap natively.Test Coverage
test_torch_sampler.py::TestApplyBadWords(11 cases): regular path (single/multi-token hit/miss, prefix in prompt) and overlap path (device-side hit/miss, host+device combined, mixed stale/fresh batches, overlapping bans → no NaN).test_torch_sampler.pysuite: 207 passed, no regressions (B200; skips are Ampere-gated).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.