[TRTLLM-14865][feat] Support occurrence penalties with beam search for TorchSampler - #17189
Conversation
zhaoyangwang-nvidia
left a comment
zhaoyangwang-nvidia
left a comment
There was a problem hiding this comment.
Reviewed the support occurrence penalties with beam search commit only (b8db3c7). I found no correctness issues.
Worth calling out what's done well: the docs were updated in the same commit; asserts vs. raises are used correctly (internal invariants asserted, the admission-time check dropped outright rather than left half-supported); and the tests compare the whole counts tensor at rtol=0, atol=0 against a host replay, so untouched slots are proven untouched — plus the e2e test's assert expected != _replay_beam_search(1.0) guards against the penalty silently not applying at all.
Five non-blocking comments inline.
Signed-off-by: Lori Ren <lorir@nvidia.com>
Signed-off-by: Lori Ren <lorir@nvidia.com>
…nalties_e2e.py The case drives a real LLM through TorchSampler, so it belongs with the other end-to-end penalty tests rather than in the op/handler unit-test module. It keeps its dummy checkpoint -- and therefore no high_cuda_memory marker, unlike its new neighbours. Leaving test_penalties.py's dynamo recompile-headroom fixture behind is a win rather than a loss: sharing a process with that module's shape-sweeping op tests was what made the case slow, and it drops from 202s to 87s. Signed-off-by: Lori Ren <lorir@nvidia.com>
9d8d9a8 to
9139366
Compare
|
@zhaoyangwang-nvidia I rebased this PR to main, pls check and review :) |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughOccurrence penalties now support beam search. Per-beam token counts are re-parented and folded before fused penalty application. Sampler wiring, tests, and sampling documentation were updated. ChangesBeam occurrence penalties
Estimated code review effort: 4 (Complex) | ~45 minutes Mergeability Score: 🔵 Low · up to This PR adds occurrence penalties to beam-search sampling, but the current tests do not cover the case where an unpenalized beam slot must be excluded from count re-parenting. The change is otherwise mergeable with explicit owner awareness of this bounded correctness risk. Sequence Diagram(s)sequenceDiagram
participant TorchSampler
participant PenaltyHandler
participant Fusions
TorchSampler->>PenaltyHandler: apply beam widths and predecessor map
PenaltyHandler->>Fusions: update beam occurrence counts
Fusions->>Fusions: re-parent and fold beam token counts
PenaltyHandler->>Fusions: apply packed occurrence penalties
Fusions-->>TorchSampler: modified logits
Possibly related PRs
Suggested reviewers: 🚥 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 (2)
tensorrt_llm/_torch/pyexecutor/sampler/penalties.py (1)
138-148: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueConsider caching the beam-index arange.
counts_rowsallocates a newtorch.arangeon every call. The rest ofPenaltyStorepersists its buffers precisely to avoid per-call allocations. The call site is the admission path, not the per-step hot path, so the cost is small. If you want the same discipline here, cache the arange once increateand reuse it.♻️ Proposed refactor
max_num_sequences: int max_beam_width: int device: torch.device + _beam_ids_cuda: torch.Tensor | None = Noneif self.max_beam_width == 1: return slots_cuda - beams = torch.arange(self.max_beam_width, device=slots_cuda.device) - return (slots_cuda.unsqueeze(1) * self.max_beam_width + beams).reshape(-1) + if self._beam_ids_cuda is None: + with torch.inference_mode(False): + self._beam_ids_cuda = torch.arange(self.max_beam_width, device=self.device) + return (slots_cuda.unsqueeze(1) * self.max_beam_width + self._beam_ids_cuda).reshape(-1)🤖 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/penalties.py` around lines 138 - 148, Cache the beam-index arange during PenaltyStore.create and store it on the instance, then update counts_rows to reuse that buffer instead of allocating torch.arange on each call. Preserve the existing single-beam fast path and ensure the cached tensor uses the appropriate device.tests/unittest/_torch/sampler/test_penalties.py (1)
652-697: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd negative coverage for the
active_cudagate.
activeis True for every slot listed inseq_slots, and the mixed-slot test setsactiveall-True as well. So theactive_cuda[seq_slots]term in botharmedandfold_oknever masks anything in this suite. That gate is what stops an unpenalized request from having its counts re-parented and folded, which matters because a beam engine calls the op for the whole batch.Include one slot in
seq_slotswithactive[slot] = Falseand assert its count rows stay unchanged.💚 Sketch of the added case
def test_beam_occurrence_counts_skip_inactive_slots() -> None: """An inactive slot in the batch must not be re-parented or folded.""" counts = torch.zeros((BEAM_SLOTS * BEAM_WIDTH, BEAM_VOCAB), dtype=torch.int32, device="cuda") active = torch.tensor([True, False], dtype=torch.bool, device="cuda").repeat(BEAM_SLOTS)[ :BEAM_SLOTS ] inactive = int((~active).nonzero()[0].item()) counts[inactive * BEAM_WIDTH + 2, 5] = 7 # only reachable via a re-parent before = counts.clone() seq_slots = torch.arange(BEAM_SLOTS, dtype=torch.int64, device="cuda") predecessor_beams = torch.full( (BEAM_SLOTS, BEAM_WIDTH), 2, dtype=torch.int32, device="cuda" ) new_tokens = torch.full((1, BEAM_SLOTS, BEAM_WIDTH), 3, dtype=torch.int32, device="cuda") Fusions.update_beam_occurrence_counts( counts, active, torch.ones(BEAM_SLOTS, dtype=torch.bool, device="cuda"), torch.ones(BEAM_SLOTS, dtype=torch.bool, device="cuda"), new_tokens, predecessor_beams, seq_slots, torch.full((BEAM_SLOTS,), BEAM_WIDTH, dtype=torch.int32, device="cuda"), BEAM_WIDTH, ) base = inactive * BEAM_WIDTH torch.testing.assert_close( counts[base : base + BEAM_WIDTH], before[base : base + BEAM_WIDTH], rtol=0, atol=0, )🤖 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 652 - 697, Add a dedicated test for update_beam_occurrence_counts that includes an inactive slot in seq_slots, seeds a count reachable only through re-parenting, and invokes the operation with representative predecessor and token data. Capture the inactive slot’s beam rows before the call and assert they remain unchanged afterward, while preserving the existing active-slot coverage. Apply the same fix in `@tests/unittest/_torch/sampler/test_penalties.py` around lines 490 - 790.
🤖 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 `@tensorrt_llm/_torch/pyexecutor/sampler/penalties.py`:
- Around line 138-148: Cache the beam-index arange during PenaltyStore.create
and store it on the instance, then update counts_rows to reuse that buffer
instead of allocating torch.arange on each call. Preserve the existing
single-beam fast path and ensure the cached tensor uses the appropriate device.
In `@tests/unittest/_torch/sampler/test_penalties.py`:
- Around line 652-697: Add a dedicated test for update_beam_occurrence_counts
that includes an inactive slot in seq_slots, seeds a count reachable only
through re-parenting, and invokes the operation with representative predecessor
and token data. Capture the inactive slot’s beam rows before the call and assert
they remain unchanged afterward, while preserving the existing active-slot
coverage.
Apply the same fix in `@tests/unittest/_torch/sampler/test_penalties.py` around
lines 490 - 790.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: ae4014e1-7c1e-4541-969f-2524fb487c8d
📒 Files selected for processing (6)
docs/source/features/sampling.mdtensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.pytensorrt_llm/_torch/pyexecutor/sampler/penalties.pytensorrt_llm/_torch/pyexecutor/sampler/sampler.pytests/unittest/_torch/sampler/test_penalties.pytests/unittest/_torch/sampler/test_penalties_e2e.py
💤 Files with no reviewable changes (1)
- docs/source/features/sampling.md
lori-ren
changed the title
zhaoyangwang-nvidia
left a comment
zhaoyangwang-nvidia
left a comment
There was a problem hiding this comment.
Approve with nits.
…m history Signed-off-by: Lori Ren <lorir@nvidia.com>
|
/bot run |
|
PR_Github #65808 [ run ] triggered by Bot. Commit: |
…er_strategy.py Signed-off-by: Lori Ren <lorir@nvidia.com>
|
PR_Github #65808 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #65816 [ run ] triggered by Bot. Commit: |
|
PR_Github #65816 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #65859 [ run ] triggered by Bot. Commit: |
…py check Signed-off-by: Lori Ren <lorir@nvidia.com>
|
/bot run |
|
PR_Github #65867 [ run ] triggered by Bot. Commit: |
|
PR_Github #65859 [ run ] completed with state |
|
PR_Github #65867 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #65903 [ run ] triggered by Bot. Commit: |
|
PR_Github #65903 [ run ] completed with state |
…r TorchSampler (NVIDIA#17189) Signed-off-by: Lori Ren <lorir@nvidia.com>
5 participants
Dev Engineer Review
PenaltyStore,PenaltyHandler, and fused operation APIs for beam metadata.QA Engineer Review
test_beam_occurrence_counts_follow_each_beam_history.test_single_beam_slot_sharing_a_beam_engine_is_routed_correctly.test_beam_search_penalties_e2e.test-db/orqa/coverage entries were provided for the new tests. Verdict: needs follow-up.Description
TorchSampler rejects repetition / presence / frequency penalties combined with beam search at admission (penalties.py:222) because PenaltyStore holds per-request token counts ([max_num_sequences]), while beam search needs them per beam. This PR extends PenaltyStore and adds repetition / presence / frequency penalties for beam search requests.
Test Coverage
test_penalties.pyto test beam search with penalties.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.