[TRTLLM-13229][feat] implement repetition / frequency / presence penalties for TorchSampler#16485
[TRTLLM-13229][feat] implement repetition / frequency / presence penalties for TorchSampler#16485lori-ren wants to merge 6 commits into
Conversation
…lties for TorchSampler Signed-off-by: Lori Ren <lorir@nvidia.com>
…kernel Signed-off-by: Lori Ren <lorir@nvidia.com>
…pler's behavior Signed-off-by: Lori Ren <lorir@nvidia.com>
Signed-off-by: Lori Ren <lorir@nvidia.com>
Signed-off-by: Lori Ren <lorir@nvidia.com>
|
/bot run |
|
PR_Github #59677 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughChangesOccurrence penalty handling
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant TorchSampler
participant OccurrencePenaltyHandler
participant TritonPenaltyKernel
participant Logits
TorchSampler->>OccurrencePenaltyHandler: Prepare request penalty state
TorchSampler->>OccurrencePenaltyHandler: Apply penalties to packed logits
OccurrencePenaltyHandler->>TritonPenaltyKernel: Launch batched occurrence kernel
TritonPenaltyKernel->>Logits: Update logits in place
TorchSampler->>OccurrencePenaltyHandler: Commit finalized token counts
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 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.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/pyexecutor/sampler/sampler.py (1)
3255-3271: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winBeam-search penalty warning is gated on sampler-level beam support, not the request's own beam width.
self._use_beam_searchreflects the sampler's configuredmax_beam_width > 1, but_is_occurrence_penalty_active(used elsewhere to actually gate penalty application) checks the request's own_get_max_beam_width(request) == 1. So for a request that itself usesbeam_width == 1inside a beam-search-capable sampler, occurrence penalties ARE applied, yet this code still emits a "penalties will be ignored" warning — misleading users about actual behavior.🐛 Proposed fix
`@override` def validate_request(self, request: LlmRequest) -> None: if self._use_beam_search: - if _has_occurrence_penalty(request): + if _get_max_beam_width(request) > 1 and _has_occurrence_penalty(request): logger.warning( "TorchSampler does not support repetition, presence, or frequency " "penalties with beam search; these penalties will be ignored." )🤖 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 `@tensorrt_llm/_torch/pyexecutor/sampler/sampler.py` around lines 3255 - 3271, Update validate_request to gate the occurrence-penalty warning on the request’s effective beam width, matching _is_occurrence_penalty_active, rather than on the sampler-level self._use_beam_search flag. Preserve the existing warning for requests whose _get_max_beam_width(request) indicates beam search, while allowing beam_width == 1 requests to apply penalties without warning.
🧹 Nitpick comments (3)
tests/unittest/_torch/sampler/test_penalties.py (1)
255-276: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd coverage for occurrence penalties with beam search.
beam_widthis exposed but every test uses1. Addtest_handler_rejects_occurrence_penalties_with_beam_searchtotests/unittest/_torch/sampler/test_penalties.py, asserting the handler’s intended validation forbeam_width=2.Coverage is otherwise strong, but beam validation is insufficient. As per path instructions, test feedback should provide a concrete file and state whether coverage is sufficient.
🤖 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_penalties.py` around lines 255 - 276, Add test_handler_rejects_occurrence_penalties_with_beam_search in test_penalties.py using _make_handler_request with beam_width=2 and assert the handler rejects occurrence penalties according to its existing validation behavior. Keep the test focused on beam search with the configured presence or frequency penalties and verify the expected exception or rejection result.Source: Path instructions
tensorrt_llm/_torch/pyexecutor/sampler/sampler.py (1)
1379-1399: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueGPU tensor fields omit the
_cudasuffix used by the sibling_FinishReasonsStore.
_FinishReasonsStorein this same file consistently suffixes device tensors (finish_reasons_cuda,max_lengths_cuda,end_ids_cuda, ...)._PenaltyStore's equivalent GPU-resident fields (repetition,presence,frequency,active,has_previous_token,counts,presence_prefix) don't follow that convention. As per coding guidelines, "use_host,_device, or_cudasuffixes when tensor location would otherwise be ambiguous."🤖 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 `@tensorrt_llm/_torch/pyexecutor/sampler/sampler.py` around lines 1379 - 1399, Rename the GPU-resident tensor fields in _PenaltyStore to use the _cuda suffix, including repetition, presence, frequency, active, has_previous_token, counts, and presence_prefix. Update every declaration, initialization, and access site consistently while preserving their existing shapes, types, and lazy None behavior.Source: Coding guidelines
tensorrt_llm/_torch/pyexecutor/sampler/ops/penalties.py (1)
194-242: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDocstring lacks Args/shape/dtype documentation, unlike its sibling.
update_occurrence_workspace(lines 40-55) documents each parameter's shape and dtype;apply_batched_occurrence_penalties_tritonis a one-liner despite having a comparably complex signature. As per coding guidelines, "Prefer docstrings for external interfaces... document public function arguments, and include tensor dimensions and constrained dtypes where applicable."🤖 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 `@tensorrt_llm/_torch/pyexecutor/sampler/ops/penalties.py` around lines 194 - 242, Expand the docstring for apply_batched_occurrence_penalties_triton to document every argument, including tensor shapes and required dtypes where applicable, and state the returned logits tensor shape/type. Match the parameter documentation style used by update_occurrence_workspace without changing the function’s behavior.Source: Coding guidelines
🤖 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 `@tensorrt_llm/_torch/pyexecutor/sampler/ops/penalties.py`:
- Around line 96-192: Separate pending-token handling from
_apply_batched_occurrence_penalties_kernel: remove its per-block read and write
of has_previous_token_ptr and do the new_tokens-to-counts update in a dedicated
single-block/preprocessing step before this kernel. Ensure the penalty kernel
reads already-updated counts without folding the pending token itself,
preserving correct behavior across all vocab blocks.
In `@tests/unittest/_torch/sampler/test_penalties_e2e.py`:
- Around line 282-335: Extend the ngram_speculative_penalties test around
_PenaltyE2ECase to enable processed log-probabilities and generation logits in
SamplingParams, then call _assert_completion_penalty_logprobs for the
speculative completion. Keep the existing reference-token equality and
speculative-stat assertions, and include test feedback stating that speculative
penalty coverage is currently insufficient.
- Around line 29-30: Annotate model_path with its Path return type, and add
explicit type annotations for both arguments and return values of the two
affected test functions. Apply the annotations to the functions near the
referenced locations while preserving their existing test behavior.
In `@tests/unittest/_torch/sampler/test_penalties.py`:
- Around line 27-31: Annotate every newly added function in
tests/unittest/_torch/sampler/test_penalties.py: add precise parameter and
return types to _col, _dense_penalty_reference, and _pack_presence_prefix, and
add -> None to each test function identified by the review. Preserve the
existing implementations and use types consistent with the surrounding test
utilities.
---
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/sampler/sampler.py`:
- Around line 3255-3271: Update validate_request to gate the occurrence-penalty
warning on the request’s effective beam width, matching
_is_occurrence_penalty_active, rather than on the sampler-level
self._use_beam_search flag. Preserve the existing warning for requests whose
_get_max_beam_width(request) indicates beam search, while allowing beam_width ==
1 requests to apply penalties without warning.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/sampler/ops/penalties.py`:
- Around line 194-242: Expand the docstring for
apply_batched_occurrence_penalties_triton to document every argument, including
tensor shapes and required dtypes where applicable, and state the returned
logits tensor shape/type. Match the parameter documentation style used by
update_occurrence_workspace without changing the function’s behavior.
In `@tensorrt_llm/_torch/pyexecutor/sampler/sampler.py`:
- Around line 1379-1399: Rename the GPU-resident tensor fields in _PenaltyStore
to use the _cuda suffix, including repetition, presence, frequency, active,
has_previous_token, counts, and presence_prefix. Update every declaration,
initialization, and access site consistently while preserving their existing
shapes, types, and lazy None behavior.
In `@tests/unittest/_torch/sampler/test_penalties.py`:
- Around line 255-276: Add
test_handler_rejects_occurrence_penalties_with_beam_search in test_penalties.py
using _make_handler_request with beam_width=2 and assert the handler rejects
occurrence penalties according to its existing validation behavior. Keep the
test focused on beam search with the configured presence or frequency penalties
and verify the expected exception or rejection result.
🪄 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: 6cfaf697-883d-41f4-aaa6-d4af6bc690ed
📒 Files selected for processing (4)
tensorrt_llm/_torch/pyexecutor/sampler/ops/penalties.pytensorrt_llm/_torch/pyexecutor/sampler/sampler.pytests/unittest/_torch/sampler/test_penalties.pytests/unittest/_torch/sampler/test_penalties_e2e.py
| @@ -0,0 +1,242 @@ | |||
| # Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. | |||
There was a problem hiding this comment.
The files under sampler/ops/ are organized by kernel source, not by feature: vanilla.py (PyTorch-native), flashinfer.py (FlashInfer-accelerated), and trtllm.py introduced in #16184 (thin wrappers over torch.ops.trtllm.*). This new penalties.py is named by feature and hosts raw @triton.jit kernels directly in this layer, which breaks both conventions.
Suggested restructuring, following the pattern established in #16184:
Move the Triton kernels and their launch logic out of sampler/ops/ and register them as custom ops in the trtllm namespace via @torch.library.custom_op("trtllm::...", mutates_args=(...)) + register_fake, following the existing trtllm::silu_and_mul pattern in tensorrt_llm/_torch/custom_ops/torch_custom_ops.py (kernel body may live in a domain module, with the registration in torch_custom_ops.py). This also matters for correctness: these kernels mutate logits / counts / has_previous_token / presence_prefix in place, and without a registered op with mutates_args, those mutations are invisible to torch.compile / graph capture.
Expose the two entry points (update_occurrence_workspace, apply_batched_occurrence_penalties_triton) from sampler/ops/trtllm.py instead, as thin keyword-only wrappers over torch.ops.trtllm.* — same style as top_p_decay_update / top_p_decay_gather there. As part of this, drop the _triton suffix from the public name (the caller shouldn't depend on the kernel implementation), and update the import in sampler.py from .ops.penalties to .ops.trtllm.
This keeps sampler/ops/ as a stable, implementation-neutral surface for the sampler, and keeps all trtllm:: op registrations discoverable in one place.
There was a problem hiding this comment.
Thanks, I agree on the naming issue. It's true that penalties.py is named by feature while the rest of sampler/ops/ is named by kernel source (vanilla.py, flashinfer.py). We may further discuss where to put these kernels.
However, on registering the kernels as trtllm::custom ops, I don't think that's necessary here:
- The dominant convention for writing Triton kernels in TRT-LLM is to use a plain Python wrapper instead of a registered op. For example, Mamba SSD (modules/mamba/) and FLA (modules/fla/) subsystems directly wraps tons of triton kernels with thin Python launchers without registering. Most kernels registered with trtllm::ops, however, are CUTLASS/cuBLAS/trtllm-gen C++ kernels instead of Triton.
- mutates_args/functionalization only matters under torch.compile, but our sampler penalty path runs eagerly (after the forward). Also, CUDA graph capture records raw Triton launches and their in-place writes fine, which is common practice for current Mamba/FLA kernels. Registration wouldn't add correctness here.
|
PR_Github #59677 [ run ] completed with state
|
Signed-off-by: Lori Ren <lorir@nvidia.com>
lori-ren
requested review from
brnguyen2,
dpitman-nvda,
schetlur-nv,
yingguo-trt and
yiqingy0
|
/bot run |
3 participants
Summary by CodeRabbit
New Features
Bug Fixes
Tests
Description
The default PyTorch TorchSampler does not apply repetition_penalty, frequency_penalty, or presence_penalty — the params are accepted by SamplingParams but silently ignored (only the TRT backend / deprecated TRTLLMSampler handle them via the C++ executor penalty kernel). These are among the most commonly used parameters in OpenAI-compatible APIs, so this is a key feature-parity gap for the PyTorch backend.
Scope:
Apply the three penalties to logits in TorchSampler, before the temperature/top-k/top-p sampling step.
Penalties depend on per-request token history: count each request's previously generated tokens, then adjust the corresponding logits — presence (penalize if token appeared), frequency (penalize proportional to count), repetition (scale appeared-token logits).
Respect existing semantics (SamplingParams definitions / C++ runtime defaults) so behavior matches the legacy backend.
Implement vectorized (no per-request Python loop), e.g. via token-count tensors, so it does not add significant host overhead and works for speculative-decoding expanded shapes.
Reference: vLLM applies all three penalties in the GPU sampler using a vectorized bincount of each request's output tokens, with no Python loop.
Test Coverage
test_penalties.pyfor validating kernel correctness and sampler behaviors.test_penalties_e2e.pyfor validating full end-to-end behaviors on TinyLlama.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.