[TRTLLM-13234][feat] Complete TorchSampler beam search: length_penalty, diversity_rate, early_stopping, VBWS, and CBA performance - #16620
Open
zhaoyangwang-nvidia wants to merge 36 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.
zhaoyangwang-nvidia
force-pushed
the
beam-length-penalty
branch
from
July 24, 2026 14:17
7167148 to
4721a74
Compare
![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
…it test The done verdict now publishes LENGTH rather than END_ID for beams that carry no reason of their own, since they end because the pool can no longer be beaten and not because they hit a token. This unit test still pinned the old value; the e2e cases already agree with the new one. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
The row_stride query added to sampler_common reads request.py_beam_width on the generation path, which MockLlmRequest and GenRequestMock do not define, so 262 cases in the A10-PyTorch-3 stage died with AttributeError before reaching what they assert. Both are single-beam; set the width to 1. These only surface on Ampere: the tests carry @force_ampere and are skipped wholesale on the Hopper machines the branch was validated on, which is why the earlier full-suite runs came back clean. Verified on an A10 -- the file now reports 372 passed, 0 failed. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
The beam-search cases in this file build a fresh engine per parametrization in one process, and Dynamo counts recompiles per code object, so the default limit of 8 is exhausted partway through and the remaining cases hard-fail under fullgraph. They only started compiling once beam search moved onto the candidate-beams-array step, which is why the file was stable before. Mirror the headroom fixture test_penalties.py already uses. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…room Dynamo counts recompiles per code object for the life of the process, and fullgraph turns exhaustion into a hard failure instead of an eager fallback. A served engine sees one shape family and never approaches the default cap of 8, but a process that builds many engines does -- a test session, or an executor worker reused across configurations. Once the cap was hit, every later beam-search request in that worker failed with "Hard failure due to fullgraph=True". Raising it from a test fixture does not reach the worker: MPI sessions only forward TRTLLM*/TLLM* environment variables, and torch._dynamo.config .recompile_limit is a plain assignment that reads no environment variable at all. Scope the higher limit to this one compiled function instead. The cap exists to catch runaway recompilation; it is not a correctness property. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
zhaoyangwang-nvidia
force-pushed
the
beam-length-penalty
branch
from
August 10, 2026 14:54
66afaaf to
5e395c6
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #65050 [ run ] triggered by Bot. Commit: |
…reason The CBA step derived its harvest mask from first_finish_reasons, which is also the request's reported finish reason and therefore has to survive for the whole request. A beam that ends on a stop word does not freeze on this path: it is pooled and its slot refills with an unrelated continuation. The reason stayed set, so the next step harvested that continuation as well -- the pool gained an entry for a hypothesis that never finished, its slot was masked again, and the request could be declared done early. Give the harvest its own one-shot signal. The finish handler raises pending_harvest for beams that finished on this step, the CBA step consumes it and lowers it in the same writeback, and admission clears it with the rest of the per-request state. test_beam_search_cba_harvest_latch_clears_after_refill drives two consecutive steps and fails on the old code with "pool grew from 1 to 2"; the existing single-step harvest test cannot reach the second step. Thanks to QiJune for spotting this and for asking that the regression test come first. NB: the latch is lowered with a tensor, not a Python scalar. Assigning False to an advanced-indexed view synchronizes, which trips the no-sync guard the beam-search step runs under. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
The PackageSanityCheck stages reject returning Any from a typed function. flashinfer.top_k is untyped, so radix_topk_op forwarded its result straight into a tuple[Tensor, Tensor] return; unpack it first. _kernel_test's wrapper had no annotations of its own, so it decayed to Any against the declared Callable[..., Any]. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
functools.wraps is untyped, so the decorated wrapper is Any and returning it violates the declared Callable[..., Any]. Annotating the wrapper itself does not help -- the decorator erases that -- so bind it to a typed name on the way out. Only the full mypy run catches this. scripts/run_mypy.sh drops to a lightweight mode with --no-warn-return-any when compiled bindings are absent, which is exactly the [no-any-return] check that fires here, so a local pre-commit on a source tree without a built wheel reports success. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…ream assert_no_cuda_sync blocks the stream with a @hostfunc and only lowers its cancel flag after the assertion. That hostfunc holds the GIL while it spins, so anything that blocks before the cancel deadlocks: the main thread waits on the stream, the stream waits on the hostfunc, and the hostfunc waits for a cancel the main thread never reaches. The speculative path runs a side-stream copier alongside the executor's own threads, which makes that race live -- the A30 stage hung here for the full 3600s pytest timeout, and the run that followed reported an out-of-bounds index from the state it left behind. Guard these two hooks with set_sync_debug_mode("error") instead. That is the property under test: torch raises on a synchronizing call. It cannot see syncs from non-torch kernels, which the stream-blocking variant would catch, but sample_async and update_requests issue none today. Left utils.util alone -- its cancel ordering is worth fixing, but not from this PR. Verified on an A30, the same GPU the stage runs on: the file goes from a 3600s hang to 7 passed in 147s, repeated three times (147.15/147.30/147.71s). Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
Two defects kept beam search from working under disaggregated serving once admission stopped rejecting it. The context server reaches the sampler with its requests already flagged as finished, so beam search finalized them. Finalization rewrites the beam rows from the beam history, and after a single context step only beam 0 has history -- every other beam is all BEAM_SEARCH_PAD_TOKEN, so set_generated_tokens gave it an empty generated sequence. The handoff reads the per-beam first token back out via getTokens().back(), which then handed the generation server the prompt tail instead of that beam's token; that side could not match it against the transferred logprob map, failed the context request, and left the generation server polling a KV transfer that never completed. Context-only requests are migrating rather than completing, so skip finalization for them and append the step's tokens instead, leaving finalization to the generation side. The disaggregated generation path also calls setup_sampler_step outside inference mode, where the beam-search store's in-place updates raise "Inplace update to inference tensor outside InferenceMode". Verified on 2xB200 with accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_beam_search (1 passed in 687s; previously hung until the KV transfer timed out). Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
zhaoyangwang-nvidia
force-pushed
the
beam-length-penalty
branch
from
August 10, 2026 15:13
5e395c6 to
238c420
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #65058 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #65050 [ run ] completed with state |
Collaborator
|
PR_Github #65058 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #65162 [ run ] triggered by Bot. Commit: |
Shixiaowei02
approved these changes
Aug 11, 2026
Collaborator
|
PR_Github #65162 [ run ] completed with state
|
A beam-search context-only request whose single sampled token is the end id produced no handoff at all: first_gen_tokens came back as None and the generation server never learned the beam finished. Such a request is finished by END_ID, which leaves it in GENERATION_COMPLETE. That state is neither isContextFinished() nor finished-due-to-length, so the branch that starts the KV transfer and calls respond_and_send_async is skipped, the request never reaches the disagg-transmission state, and llmRequest.cpp never builds ContextPhaseParams -- dropping the tokens, the scores and the logprob map together. This is the completion loss the admission-time rejection used to describe. Mask the end id to the "no end token" sentinel for these requests at admission instead. The finish handler then records no END_ID finish, so the request completes its context phase and hands off normally, carrying the end token as first_gen_tokens for the generation side to pool via the existing pending_harvest latch. The CBA op reads the same masked value out of the store, which also keeps the end candidate in its beam slot -- so the per-step mask the sampler used to build for it, a clone plus two tolist() calls on every beam-search step, is gone. Single-beam disaggregation keeps its current end-id behaviour. Adds a test that picks the end id after a probe run so the context step finishes deterministically on its first and only token; it fails with "context phase produced no first_gen_tokens" without this fix. Verified on 2xB200: the new test passes 3/3 (and fails when the fix is reverted), tests/unittest/_torch/sampler is 1136 passed with one pre-existing VBWS C++/Python mismatch, and accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_beam_search passes in 773s. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #65222 [ run ] triggered by Bot. Commit: |
4 tasks
Collaborator
Author
Collaborator
|
PR_Github #65222 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #65426 [ run ] triggered by Bot. Commit: |
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
6 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: every mode runs on a candidate-beams array(CBA) mirroring the C++ kernels — end-token candidates enter a best-K
finished pool and the vacated slot is refilled, so one lineage can
contribute several hypotheses. The mode only selects the done verdict:
1stops once the pool holds
best_ofcandidates,0and2additionallyrequire that no active beam can still beat the pool (differing in how
optimistic that attainability bound is).
detection, next-step harvest into the pool, slot refill; the stop-word
window follows beam swaps (also fixes pre-existing swap-unaware matching).
logprobs(C++logProbsTiled/logProbsCBAanalogs: per-steplog-prob history + pool snapshots).
beam_width_array), including fixes for severalpre-existing bugs that made VBWS unusable: an out-of-bounds width lookup
past the array end (C++
getBeamWidthByIterclamped with the globalcapacity constant rather than the array's own length, which starved the
request in the micro-batch scheduler and hung decoding), a bad full-width
view in finished-state propagation, generation logits offset by the
per-iteration width while the forward path lays rows out at the static
admission width, and the padding sentinel leaking into request token
histories.
Behavior changes
early_stopping=1no longer freezes a finished beam in its slot. Itpreviously stopped once every slot held a finished beam, which a comment
claimed was equivalent to the C++ verdict. It is not: C++ counts
hypotheses accumulated in the pool, and because a finished candidate
vacates its slot there, one lineage can contribute several. vLLM does the
same. Beam contents for this mode can therefore differ from before.
length_penaltyandearly_stoppingnow defaultto
nulland defer to the engine (0.0/1), where the schemapreviously restated
1.0/false. A beam-search request that never setlength_penaltywas normalizing scores by length and now ranks by the rawcumulative log-probability; set
"length_penalty": 1.0to keep the oldranking.
early_stoppingacceptsfalse/true/"never"and rejectsintegers outside that set instead of silently reinterpreting them.
pool the context server builds is not part of the handoff
(
ContextPhaseParamscarries only the first generated tokens) and iscleared on the generation side, so a completion found during the context
phase would be silently dropped. This now applies to both samplers: the
C++ decoder maintains the pool unconditionally, so TRTLLMSampler was
affected too, and the previous check only rejected the exhaustive modes.
This costs the beam-score handoff added earlier in this branch; the
seeding code is kept and annotated, and becomes live again once the
handoff carries the pool (TRTLLM-14792).
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. Under
early_stopping1and0, three of four beams match token for token;the fourth differs by the trailing-EOS convention, which changes a short
finished hypothesis's normalized score and hence its rank.
configurations; the new features add no measurable overhead when enabled,
zero overhead when disabled. NB: measured before
early_stopping=1movedonto the CBA path, so that mode's numbers predate the change and want a
re-run.
best_of > 1; a decreasingbeam_width_array(only non-decreasingschedules are defined, and finalize reads the array maximum); a
best_ofthat differs from
max_beam_width.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.