[TRTLLM-13234][feat] Complete TorchSampler beam search: length_penalty, diversity_rate, early_stopping, VBWS, and CBA performance - #16620
Open
zhaoyangwang-nvidia wants to merge 19 commits into
Conversation
zhaoyangwang-nvidia
force-pushed
the
beam-length-penalty
branch
3 times, most recently
from
July 21, 2026 06:54
a0cc88a to
bcc284f
Compare
zhaoyangwang-nvidia
changed the title
zhaoyangwang-nvidia
force-pushed
the
beam-length-penalty
branch
2 times, most recently
from
July 21, 2026 09:36
d246f47 to
f11c23f
Compare
zhaoyangwang-nvidia
changed the title
zhaoyangwang-nvidia
changed the title
zhaoyangwang-nvidia
force-pushed
the
beam-length-penalty
branch
7 times, most recently
from
July 22, 2026 10:51
d5dab0e to
7167148
Compare
zhaoyangwang-nvidia
changed the title
zhaoyangwang-nvidia
marked this pull request as ready for review
July 22, 2026 10:54
zhaoyangwang-nvidia
requested review from
JunyiXu-nv,
YihuiLu512,
arysef and
laikhtewari
July 22, 2026 10:54
Contributor
|
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:
WalkthroughBeam-search sampling now supports length penalties, diversity rates, and configurable early stopping. Candidate selection uses staged top-k dispatch, while non-default early-stopping modes use CBA state and compiled step math. Request handling, beam stores, history finalization, documentation, and tests were updated. ChangesBeam Search Sampling
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Suggested labels: Suggested reviewers: Sequence Diagram(s)sequenceDiagram
participant Request
participant TorchSampler
participant BeamSearchMixin
participant BeamSearchKernel
participant BeamHistoryBuilder
Request->>TorchSampler: provide beam-search parameters
TorchSampler->>BeamSearchMixin: resolve beam strategy
BeamSearchMixin->>BeamSearchKernel: dispatch beam-search kernel
BeamSearchKernel->>TorchSampler: update beams and CBA state
TorchSampler->>BeamHistoryBuilder: provide CBA snapshots
BeamHistoryBuilder->>Request: return ranked beam histories
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
ixlmar
reviewed
Jul 23, 2026
ixlmar
left a comment
ixlmar
left a comment
Collaborator
There was a problem hiding this comment.
Many of the comments currently reference details of the C++ code. I think we should try to make these self-contained, such that the new code remains maintainable even if the C++ bits get removed at some point.
| @@ -849,6 +1475,8 @@ def make_metadata(seq_len: int) -> BeamSearchMetadata: | |||
| dtype=torch.int32), | |||
| seq_offsets=seq_offsets, | |||
| beam_idx_arange=beam_idx_arange, | |||
| beam_gen_lengths=torch.zeros((max_batch_size, beam_width), | |||
Collaborator
There was a problem hiding this comment.
Since the _cba beam-search variant is currently implemented in a code path separate from the regular beam search, all the tests would need to be adapted to cover those two paths.
Generally, I find it difficult to see how much beam_search_sampling_batch_cba and beam_search_sampling_batch differ from each other. Superficially, I would expect beam_search_sampling_batch_cba to be a superset of beam_search_sampling_batch, such that there should be potential for code reuse, which would improve both readability and testability. @stnie Would you perhaps have time to look over these changes and comment on that point, in particular?
Collaborator
Author
There was a problem hiding this comment.
Thanks — I've taken this in two parts:
Testability: both paths are now covered independently. The CBA path has dedicated unit tests (test_beam_search_cba_*), including the early_stopping=NEVER max_seq_len done-bound and a negative length_penalty, alongside the regular-path tests.
Code reuse: I extracted the shared front-end (reshape, temperature, cache-indirection snapshot, log_softmax) into _beam_step_preprocess, called by both ops. Beyond that, the two aren't in a superset relationship: the finished-beam handling differs (regular keeps the EOS column; CBA masks the row to harvest it), and the tail is a genuinely different algorithm — the regular op updates the store straight from the top-k, while the CBA op maintains the candidate-beams array, computes the attainability done-verdict, and merges the pool via the torch.compiled _cba_step_math.
…lidate early_stopping in TorchSampler Implement length_penalty and beam_search_diversity_rate for the PyTorch sampler's beam search, matching the C++ decoder semantics (applyLengthPenalty and the diversityRate beam-slot top-k adjustment), and reject unsupported early_stopping modes instead of silently ignoring them: - Rank beam candidates by (cum_log_prob + diversity_rate * source_beam_index) / gen_len**length_penalty via a two-stage top-k (per-beam top-k on raw scores, then adjust only the bw_in * bw_out survivors) that is mathematically equivalent to adjusting the full candidate matrix but never touches the vocab axis. Stored cum_log_probs remain raw. - Track per-beam generated lengths (frozen when a beam finishes) in a new BeamSearchStore.beam_gen_lengths buffer, only maintained when length_penalty is active. - Use flashinfer's radix top-k for the beam candidate selection (~5x faster than torch.topk on vocab-sized rows); vanilla.py stays flashinfer-free via a topk_fn injection point. - Reject early_stopping values other than 1 in validate_request instead of silently ignoring them (TorchSampler's stop criterion is equivalent to early_stopping=1; the exhaustive modes need a separate finished-candidate pool). - Document the new parameters in docs/source/features/sampling.md and fix the beam_search_diversity_rate docstring in sampling_params.py (it was a copy of repetition_penalty's text with a wrong default). Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…orchSampler Implement the early_stopping != 1 modes via a candidate-beams array (CBA), mirroring the C++ beamSearchKernels + stopCriteria workflow: - beam_search_sampling_batch_cba (ops/vanilla.py): selects the top 2 * beam_width expansion candidates by raw score (+ diversity), routes end-token candidates (rank < beam_width) into the CBA — which keeps the best beam_width finished paths seen so far by length-normalized score, with paths snapshotted eagerly through the cache indirection — and continues all beam slots with the best active candidates, so exploration never narrows. - Stop words (any length, following the C++ stop-criteria split): the finish handler latches per-beam stop-word finishes after each step; at the start of the next step the op harvests the latched beams into the CBA and masks their rows so the freed slots refill with active candidates — the torch equivalent of the C++ kernel coercing a finished beam into a top-ranked end-token candidate. The op also reorders the finish handler's rolling stop-word window by the step's predecessor beams so multi-token matching stays correct across beam swaps. - A request is done when the CBA is full and the best candidate's attainable normalized score cannot beat the worst CBA entry (early_stopping == 0 bounds with the current length; other values with max_seq_len when length_penalty > 0, following the C++ kernel). The verdict is published by marking all beam slots finished, driving the regular stop machinery. - Finalization merges the CBA with the current active paths and returns the top beam_width by normalized score (_prepare_beam_history_cba). - logprobs and variable beam width are rejected for these modes (validate_request). - early_stopping joins the beam-search strategy tuple and grouping key so CBA and frozen-slot groups stay uniform; early_stopping == 1 keeps the existing (equivalent) frozen-slot path unchanged. Verified against TRTLLMSampler: top-beam parity on 8 (length_penalty, diversity_rate, early_stopping) combinations; stop-word cases match TorchSampler's existing early_stopping == 1 behavior. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
- Move the flashinfer-vs-torch.topk row-width dispatch (topk_op) from ops/flashinfer.py (now a pure wrapper, radix_topk_op) to sampling_utils, the layer that selects implementations; route every beam-search top-k in ops/vanilla.py through the injected topk_fn. - Use the two-stage selection also on the no-adjustment path and for the CBA candidate selection (both are calls into beam_candidate_topk now, with the tiny-vocab clamps absorbed into it): equivalent to the flat top-k, single code path, and faster with the radix backend at typical beam-search batch sizes (bs=32/K=4: 100us vs 152us; crossover to flat around bs*K ~ 512 rows, +27us there — invisible behind the host-bound step time either way). Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…arch The finish handler's rolling stop-word window is indexed by beam slot, but a slot's occupant changes on beam swaps, so multi-token stop-word matching could compare a new occupant's tokens against its predecessor's history (the BeamSearchStore.predecessor_beams docstring already described this need, but nothing consumed it). Reorder the window by the step's predecessor beams in the beam search op — same fix the CBA path already carries — and cover both paths with unit tests. Note: the remaining torch-vs-TRTLLMSampler difference on stop-word- terminated beams is an output-semantics one (TorchSampler excludes the matched stop-word tokens from token_ids, TRTLLMSampler keeps them, which also shifts the raw-score ranking of such short beams); the search behavior itself is aligned. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…ing modes Mirror the C++ logProbsTiled / logProbsCBA workspaces: the CBA path records each step's per-slot sampled log-prob in a device history buffer (BeamSearchStore.original_log_probs, the log-prob analog of original_tokens), snapshots per-token log-probs into the CBA alongside the token paths (end-token/harvested insertions), and merges them at finalization so every returned beam carries its sampled-token log-probs and raw cumulative log-prob. Drops the corresponding validate_request rejection; only variable beam width remains unsupported with early_stopping != 1. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…n TorchSampler Variable-Beam-Width-Search (beam_width_array) crashed on the Torch sampler; make it work on both the default and the exhaustive early_stopping paths: - LlmRequest.get_beam_width_by_iter (python override): the C++ binding clamps the decoding-iteration index with the global kMaxBeamWidthArrayLength constant (assuming a padded array) and reads past the end of the raw user array once decoding outlives it, returning garbage widths (observed: beam width 0 crashing the top-k kernel). Clamp with the actual array length. - beam_search_sampling_batch: the finished-state propagation viewed the [batch, beam_width_out] gather as the full store width, crashing on any step with beam_width_out != max_beam_width; write only the produced columns. - Both beam ops pad returned tokens to the store beam width (the batched sampling buffers are allocated at max_beam_width); the padding is never consumed since finalization rewrites beam tokens. - CBA path: per-request pool capacity (cba_caps == the request's maximum beam width, C++ nBM) for admission rank, merge, and the done bound; candidate selection widened to beam_width_out + max(beam_width_out, beam_width_in) so narrowing steps keep enough active candidates; finalization separates pool width, current active width, and output width. Verified against TRTLLMSampler: top-beam parity with beam_width_array=[2, 3, 4] under early_stopping 1 and 0. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
The exhaustive early_stopping path cost ~3.4ms/step over the default path at bs=32/beam=4; profiling attributed it to ~90 small-kernel launches in the step op (half the op wall was launch gaps) and to per-request host D2H staging in the finalize preparation (~8 copies x batch requests per step, host-call-count bound): - Extract the step's small-tensor math into _cba_step_math and run it under torch.compile (fullgraph, mark_dynamic on the batch and snapshot-length dims, following the Fusions pattern); the data-dependent masked_select slot selection is replaced by a scatter through a (K+1)-wide buffer, and the length penalty becomes a per-request exponent tensor (0 == off) so the graph is branch-free. CPU callers keep the eager function. Store writebacks stay eager (mixed advanced/basic indexing). - Snapshot and pool-merge widths are bounded by the host-known maximum generated length (BeamSearchMetadata.max_gen_len) instead of the full max_seq_len. - _prepare_cba_group_host batches the finalize D2H state for all CBA-mode requests of a step into one copy per tensor (plus a batched stop-criteria verdict); the per-request builders slice host rows. Step op overhead over the default path: ~1000us -> ~180us; end-to-end beyond-forward for early_stopping=0 (sampler microbench, H200, bs=32, beam=4, gpu-busy 2000us): ~5.9ms -> ~1.1ms per step. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
zhaoyangwang-nvidia
force-pushed
the
beam-length-penalty
branch
from
July 24, 2026 14:17
7167148 to
4721a74
Compare
Non-functional review follow-ups, no runtime behavior change: - Document source_beam_index in the beam_search_diversity_rate description (docs/source/features/sampling.md and sampling_params.py): it is the rank of the source beam among the step's input beams, ordered by cumulative log-probability (0 = strongest). - Note in _cba_step_math's docstring that it is torch.compile'd, so it must stay free of data-dependent shapes and in-place ops on its inputs (a reviewer nearly added masked_fill_ suggestions without this context). - Fold the snap_arange mark_dynamic into the preceding loop tuple. - Fix the _prepare_cba_group_host NVTX range name (was a stale "maybe_create_beam_histories"). Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
Non-functional review follow-ups on the CBA beam-search path: - Group the CBA-only fields of BeamSearchMetadata into a dedicated CBAState dataclass, so callers gate on a single `cba is not None` check instead of testing each field. stop_past_tokens stays top-level (used by the regular path too). - Promote the dict returned by _prepare_cba_group_host to a CBAGroupHost dataclass; consumers use attribute access. The stricter typing surfaced a latent Optional[int] seq-slot, now asserted. - Pass _cba_step_math's 25 arguments by name at the call site (verified the torch.compile'd wrapper accepts kwargs with no extra recompilation and identical results). - Use cast(int, ...) instead of int(...) for the length-derived width. - Expand the "at least K non-end candidates" comment to explain why the invariant holds under stop words and multi-EOS models. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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/sampler.py`:
- Around line 4022-4042: The CBAState construction in the beam-search strategy
group must only occur when the group’s requests use CBA. Use
`_request_uses_cba(...)` to gate the existing `CBAState` value and pass
`cba=None` for non-CBA groups, preserving all current CBAState fields and
calculations.
- Around line 4161-4170: Update the CBAGroupHost construction to populate
should_stop from d2h_copier(store.batch_dones[slots_cuda]) instead of the
per-beam finish-reason-derived should_stop value. Preserve the existing slot
mapping and other snapshot fields unchanged.
🪄 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: e5bc2ad3-df0c-4f99-8ee0-03dea14c3496
📒 Files selected for processing (4)
tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.pytensorrt_llm/_torch/pyexecutor/sampler/sampler.pytensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.pytests/unittest/_torch/sampler/test_beam_search.py
🚧 Files skipped from review as they are similar to previous changes (3)
- tests/unittest/_torch/sampler/test_beam_search.py
- tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py
- tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py
Promote two positional tuples to NamedTuples so their fields are self-documenting, without changing behavior: - _cba_step_math now returns a CBAStepResult NamedTuple (not a dataclass: the function is torch.compile'd with fullgraph=True, and a NamedTuple is a valid graph output while a custom dataclass risks a graph break; verified the compiled path stays fullgraph-clean with identical results). - The BeamSearch strategy tuple becomes a NamedTuple; match/case sequence patterns and indexing are unchanged. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
Address review: comments referenced C++ kernel symbols (beamSearchKernelsTemplate.h, applyLengthPenalty, nBM, logProbsTiled, logProbsCBA, batchDones, the finalize kernel) to explain the torch implementation. Since that C++ code may be removed, rewrite the doc-strings and comments to describe the algorithm in their own terms instead of delegating to C++ symbol names. No behavior change. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
Move the beam-search kernels and their metadata out of ops.vanilla into a dedicated sampler/beam_search.py (candidate top-k, the regular and CBA exhaustive-early-stopping paths, the torch.compile'd CBA step, and the BeamSearchMetadata / CBAState / CBAStepResult dataclasses). ops.vanilla keeps only the non-beam sampling kernels (greedy, top-k/top-p, rejection, Fusions). sampling_utils re-exports the moved symbols so downstream imports (sampler, tests) are unchanged in spirit; test imports are redirected to the new module. Pure move — verified byte-for-byte parity of every store tensor the regular and CBA paths write back, before and after. This addresses the review request to make the beam-search code easier to read and test by giving it its own home, without merging the regular and CBA paths (which diverge in candidate handling and the CBA-only pool/compile machinery). Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
Replace the raw int early_stopping with a named IntEnum (BeamSearchEarlyStop) that mirrors HF's tri-state values TRUE / FALSE / NEVER, with docstrings documenting each mode's stopping bound. Addresses a beam-search review comment. The top-level SamplingParams API stays Optional[int]; the enum is IntEnum so it remains interchangeable with the raw config integers and the strategy grouping key. Also reword the early_stopping / length_penalty docstrings per review: the exhaustive-mode bound is stated in terms of the beam's current score (scores decrease monotonically with length) versus no upper bound for unfinished beams when length_penalty > 0. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
TorchSampler allocates beam buffers at max_beam_width and slices to the per-request width, so a request may use a beam width smaller than the engine's max_beam_width. Relax the executor's beam-width validation accordingly: reject only widths above max_beam_width; permit smaller widths for TorchSampler, except together with log-probs (not yet wired through the variable-beam-width path). The legacy TRT-LLM sampler still requires beam_width == max_beam_width. Update the parameter-validation test to cover both samplers: the legacy sampler still errors, while TorchSampler returns best_of beams for a smaller width and errors only when log-probs are also requested. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…earch module Move the pure beam-search data structures out of the TorchSampler class body and the module top level into the beam_search module, next to the CBA/metadata dataclasses they belong with: BeamSearchStore, BeamHistory, _BeamHistoryLogProbsSlices, _BeamHistoryTensors, and the _gather_beam_path tensor helper. They are plain dataclasses / pure tensor functions with no TorchSampler dependency, so this is a mechanical relocation; sampling_utils re-exports them and sampler.py imports from there. No behavior change. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
Replace the BeamSearchMixin + BeamSearchWithProbs/BeamSearchSampleOnly trio with a BeamSearchStep base and two subclasses, RegularBeamSearchStep (early_stopping == TRUE) and CBABeamSearchStep (FALSE / NEVER). The stopping mode, previously an in-sample() if-branch selecting one of two free functions, is now the _select_and_update hook each subclass implements; shared temperature preprocessing stays in the base sample(). With-probs vs. sample-only is no longer a subclass for beam search (it carried no separate sampling path, only a return_probs flag): it becomes a constructor flag the dispatcher injects via with_computes_probs(), so beam search no longer needs the WithProbs/SampleOnly variants. Dispatch selects Regular/CBA from the early_stopping already in the grouping key. Other strategies are unchanged. No behavior change (verified equivalent across all early_stopping x return_probs combinations). Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
- Deduplicate the CBA tensor fields shared by BeamSearchStore and CBAState into a _CBAFields base class, documenting each field's shape/semantics once instead of in both classes. - Trim the design-rationale prose from the BeamSearchEarlyStop members, BeamSearchStep, and BeamHistory docstrings down to what each field/mode does. - Replace the injected topk_fn with an internal _beam_topk dispatcher in the beam_search module that calls flashinfer's radix top-k directly on wide rows and torch.topk on small ones; drop the now-unused topk_op. flashinfer is already a hard TorchSampler dependency, and its ops module stays import-safe without it. - Beam-search kernel unit tests now run on CUDA (via a _kernel_test decorator), cover early_stopping=NEVER's max_seq_len done bound and a negative length_penalty, assert finished-beam markers actually move with the predecessor swap, and drop a duplicated use_beam_search test. Tighten the smaller-beam-width rejection match string. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
The regular and CBA beam-search step ops repeated the same front-end (reshape, temperature, cache-indirection buffer snapshot, log_softmax). Factor it into _beam_step_preprocess so both call it once; no behavior change. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
- Use maybe_mark_dynamic for snap_arange in the CBA step: newer dynamo raises ConstraintViolationError when a mark_dynamic dim is specialized to a constant, which the bounded snapshot width is; allow specialization. - Drop the assert_no_cuda_sync guard in test_beam_search_sampling_batch_basic: now that the kernel unit tests run on GPU, ordinary step ops (softmax, scalar-float temperature division) synchronize under recent torch, so the guard's no-sync contract no longer holds and is unrelated to the logic under test. Verified on GPU: 311 passed, 245 skipped (model-weight e2e) for the full test_beam_search.py. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
… TorchSampler
Implement length_penalty and beam_search_diversity_rate for the PyTorch sampler's beam search, matching the C++ decoder semantics (applyLengthPenalty and the diversityRate beam-slot top-k adjustment):
Summary by CodeRabbit
New Features
TorchSamplerbeam-search support aligned with C++ semantics, includinglength_penalty,beam_search_diversity_rate, exhaustiveearly_stoppingmodes, variable beam widths, stop-words handling, andlogprobs.beam_candidate_topk) with atopk_opdispatcher that uses FlashInfer radix top-k when beneficial (viaradix_topk_op) and otherwise falls back totorch.topk.Documentation
docs/source/features/sampling.mdwithSamplingParamsdescriptions forlength_penalty,beam_search_diversity_rate, andearly_stopping, including stop-condition behavior and notes about default vs exhaustive/heuristic modes and cumulative-logprob normalization.Code/Behavior Updates
LlmRequest.get_beam_width_by_iter()override added to safely derive iteration-dependent beam widths frombeam_width_array, including shape normalization and clamping to prevent out-of-bounds/garbage-width behavior.torch.compile-friendly refactor for the exhaustive CBA step math (with CUDA/eager branching), plus explicit eager writebacks/finalization transfers to preserve snapshot/logprob semantics safely.early_stoppingvalues.Tests
tests/unittest/_torch/sampler/test_beam_search.pyforlength_penaltyranking semantics, diversity effects, FlashInfertopk_opparity (including-inf-dominated rows), and extensive CBA/EOS/slot/stop-word window behaviors.beam_gen_lengthsfor beam-search/CBA paths.Dev Engineer Review
Correctness & Semantics
length_penalty,beam_search_diversity_rate, andearly_stoppingthrough TorchSampler grouped strategy selection into the beam-search kernels.beam_width_arrayshapes and clamping derived widths to the valid configured range.beam_candidate_topk) and incorporateslength_penalty/diversity into selection keys while preserving raw (unnormalized) storedcum_log_probs, matching intended semantics and covered by unit tests.stop_past_tokens/ CBA reordering logic), and_cba_step_math) with explicit eager writebacks to ensure correct tensor snapshot/logprob handling.next_tokens, relying on candidate selection to return valid token ids.Performance/Implementation
radix_topk_op) and atopk_opdispatcher to reduce overhead while maintaining parity withtorch.topk.CBAState) and grouped host snapshot plumbing (CBAGroupHost) to support batched CBA finalization.API/Code Consistency
BeamSearchMetadata,BeamSearchStore) with new CBA tensors (generated lengths, CBA token/score/logprob/length/capacity buffers, termination verdicts, etc.) and updated call signatures to thread additional sizing inputs (prompt_lens_cuda,beam_caps_cuda).Risk Areas to Re-check
early_stopping != 1vsearly_stopping == 1).QA Engineer Review
Test files touched
tests/unittest/_torch/sampler/test_beam_search.pyTest functions added/modified (as present in file)
test_beam_search_e2etest_beam_search_disagg_e2etest_beam_search_large_beam_width_regressiontest_beam_search_sampling_batch_basictest_beam_search_sampling_batch_length_penaltytest_beam_candidate_topk_equivalencetest_topk_op_beam_paritytest_beam_search_sampling_batch_diversity_ratetest_beam_search_cba_insert_and_slotstest_beam_search_cba_done_when_unbeatabletest_beam_search_cba_replace_mintest_beam_search_cba_harvest_stop_word_beamtest_beam_search_cba_reorders_stop_windowtest_beam_search_sampling_batch_reorders_stop_windowtest_beam_search_sampling_batch_disagg_handofftest_create_beam_historytest_finish_beamsclass TestParameterValidation:test_use_beam_search_falsetest_use_beam_search_ommittedtest_smaller_beam_widthtest_logprobs_trtllm_samplertest_logprobs_torch_samplerCoverage mapping
tests/unittest/...(this PR does not modifytests/integration/test_lists//test-db//qa/entries).Verdict
Description
TorchSampler silently ignored most beam-search sampling parameters (only the
deprecated TRTLLMSampler honored them). This PR completes the support, matching
the C++ decoder semantics, and optimizes the new paths.
New support
length_penalty: candidates ranked bycum_log_prob / gen_len**penalty(matching C++
applyLengthPenalty); returnedcum_log_probsstay raw.beam_search_diversity_rate:rate * source_beam_indexadded to theranking score (matching the C++ beam-slot top-k).
early_stopping, all modes:1(default) maps to the existingfrozen-slot path; the exhaustive modes (
0/other) get a candidate-beamsarray (CBA) mirroring the C++ kernels — end-token candidates enter a
best-K finished pool, all slots continue exploring, and a request stops
when no active beam can beat the pool (per-mode attainability bounds).
post-step detection, next-step harvest into the pool, slot refill; the
stop-word window now follows beam swaps (also fixes the pre-existing
swap-unaware matching on the default path).
logprobswith exhaustive modes (C++logProbsTiled/logProbsCBAanalogs: per-step log-prob history + pool snapshots).
beam_width_array) with all modes; fixes twopre-existing bugs that crashed any VBWS request on the Torch sampler
(out-of-bounds width lookup past the array end; a bad full-width view in
finished-state propagation).
Optimizations
beam_candidate_topk): per-beam top-k on raw scores,then adjust only the
bw_in x bw_outsurvivors — equivalent to adjustingthe full candidate matrix, without touching the vocab axis.
torch.topk on vocab-sized rows), behind a size-dispatching
topk_op;ops/vanilla.pystays flashinfer-free via atopk_fninjection point.the
Fusionspattern) and group-batched finalize D2H (one copy pertensor per step instead of ~8 per request): exhaustive-mode overhead went
from ~5.9ms to ~1.1ms per step (sampler microbench, H200, bs=32, beam=4).
Results
diversity_rate, early_stopping, stop/VBWS) e2e combinations.
configurations; the new features add no measurable overhead when enabled,
zero overhead when disabled.
output-convention differences vs TRTLLMSampler (trailing EOS and matched
stop-word tokens in
token_ids) are documented in the code.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.