[TRTLLM-13234][feat] Complete TorchSampler beam search: length_penalty, diversity_rate, early_stopping, VBWS, and CBA performance by zhaoyangwang-nvidia · Pull Request #16620 · NVIDIA/TensorRT-LLM

zhaoyangwang-nvidia · 2026-07-20T10:53:26Z

… 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):

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).

Summary by CodeRabbit

New Features
- Implemented complete PyTorch TorchSampler beam-search support aligned with C++ semantics, including length_penalty, beam_search_diversity_rate, exhaustive early_stopping modes, variable beam widths, stop-words handling, and logprobs.
- Added generated-length tracking and diversity-based beam-slot ranking for correct beam ordering and candidate selection.
- Introduced CBA (candidate-beam array) support for exhaustive early stopping: finished-candidate pools, EOS/slot handling, and stop-word window reordering across beam swaps.
- Implemented two-stage beam candidate selection (beam_candidate_topk) with a topk_op dispatcher that uses FlashInfer radix top-k when beneficial (via radix_topk_op) and otherwise falls back to torch.topk.
Documentation
- Extended docs/source/features/sampling.md with SamplingParams descriptions for length_penalty, beam_search_diversity_rate, and early_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 from beam_width_array, including shape normalization and clamping to prevent out-of-bounds/garbage-width behavior.
- TorchSampler beam-search plumbing extended to propagate beam-search strategy parameters, allocate new CBA buffers, and route beam-history preparation through a CBA-specific builder when exhaustive early stopping is enabled.
- Added 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.
- Added validation/dispatch wiring for unsupported early_stopping values.
Tests
- Expanded/updated unit tests in tests/unittest/_torch/sampler/test_beam_search.py for length_penalty ranking semantics, diversity effects, FlashInfer topk_op parity (including -inf-dominated rows), and extensive CBA/EOS/slot/stop-word window behaviors.
- Updated test metadata builders to include required beam_gen_lengths for beam-search/CBA paths.

Dev Engineer Review

Correctness & Semantics
- Strategy/parameter resolution now propagates length_penalty, beam_search_diversity_rate, and early_stopping through TorchSampler grouped strategy selection into the beam-search kernels.
- Beam-width iteration handling is made robust by normalizing potentially nested beam_width_array shapes and clamping derived widths to the valid configured range.
- Beam candidate selection now uses a two-stage top-k flow (beam_candidate_topk) and incorporates length_penalty/diversity into selection keys while preserving raw (unnormalized) stored cum_log_probs, matching intended semantics and covered by unit tests.
- Stop-word history is kept consistent across beam swaps/reordering for both:
  - the exhaustive CBA path (via stop_past_tokens / CBA reordering logic), and
  - the ES=1 path (via predecessor-based window reordering).
- Exhaustive early stopping now uses CBA-specific state/buffers plus a compiled-step math refactor (_cba_step_math) with explicit eager writebacks to ensure correct tensor snapshot/logprob handling.
- Removed the prior vocab-modulo adjustment for next_tokens, relying on candidate selection to return valid token ids.
Performance/Implementation
- Added FlashInfer radix top-k wrapper (radix_topk_op) and a topk_op dispatcher to reduce overhead while maintaining parity with torch.topk.
- Candidate selection and exhaustive finalization paths were refactored to reduce sampler overhead (two-stage top-k; batched CBA finalization with careful tensor writes).
- Added structured CBA state (CBAState) and grouped host snapshot plumbing (CBAGroupHost) to support batched CBA finalization.
API/Code Consistency
- Extended internal dataclasses/metadata (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 dispatch correctness and completeness across CBA vs non-CBA paths (early_stopping != 1 vs early_stopping == 1).
- Variable beam width per iteration correctness (beam-width array shape handling + clamping).
- Correct alignment of stop-word window ordering with predecessor-beam mappings during beam swaps (especially in CBA updates).

QA Engineer Review

Test files touched
- tests/unittest/_torch/sampler/test_beam_search.py
Test functions added/modified (as present in file)
- test_beam_search_e2e
- test_beam_search_disagg_e2e
- test_beam_search_large_beam_width_regression
- test_beam_search_sampling_batch_basic
- test_beam_search_sampling_batch_length_penalty
- test_beam_candidate_topk_equivalence
- test_topk_op_beam_parity
- test_beam_search_sampling_batch_diversity_rate
- test_beam_search_cba_insert_and_slots
- test_beam_search_cba_done_when_unbeatable
- test_beam_search_cba_replace_min
- test_beam_search_cba_harvest_stop_word_beam
- test_beam_search_cba_reorders_stop_window
- test_beam_search_sampling_batch_reorders_stop_window
- test_beam_search_sampling_batch_disagg_handoff
- test_create_beam_history
- test_finish_beams
- class TestParameterValidation:
  - test_use_beam_search_false
  - test_use_beam_search_ommitted
  - test_smaller_beam_width
  - test_logprobs_trtllm_sampler
  - test_logprobs_torch_sampler
Coverage mapping
- Covered by unit tests in tests/unittest/... (this PR does not modify tests/integration/test_lists/ / test-db/ / qa/ entries).
Verdict
- sufficient (broad semantic coverage for ranking/selection effects, FlashInfer-vs-torch top-k parity, and detailed CBA/EOS/slot and stop-word reordering edge cases).

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 by cum_log_prob / gen_len**penalty
(matching C++ applyLengthPenalty); returned cum_log_probs stay raw.
beam_search_diversity_rate: rate * source_beam_index added to the
ranking 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: 1
stops once the pool holds best_of candidates, 0 and 2 additionally
require that no active beam can still beat the pool (differing in how
optimistic that attainability bound is).
Stop words (any length) via the C++ stop-criteria split: post-step
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/logProbsCBA analogs: per-step
log-prob history + pool snapshots).
Variable beam width (beam_width_array), including fixes for several
pre-existing bugs that made VBWS unusable: an out-of-bounds width lookup
past the array end (C++ getBeamWidthByIter clamped with the global
capacity 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=1 no longer freezes a finished beam in its slot. It
previously 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.
HTTP schema defaults. length_penalty and early_stopping now default
to null and defer to the engine (0.0 / 1), where the schema
previously restated 1.0 / false. A beam-search request that never set
length_penalty was normalizing scores by length and now ranks by the raw
cumulative log-probability; set "length_penalty": 1.0 to keep the old
ranking. early_stopping accepts false / true / "never" and rejects
integers outside that set instead of silently reinterpreting them.
Beam search is rejected under disaggregated serving. The finished
pool the context server builds is not part of the handoff
(ContextPhaseParams carries only the first generated tokens) and is
cleared 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

Two-stage top-k (beam_candidate_topk): per-beam top-k on raw scores,
then adjust only the bw_in x bw_out survivors — equivalent to adjusting
the full candidate matrix, without touching the vocab axis.
flashinfer radix top-k for all beam-search top-k sites (~5x faster than
torch.topk on vocab-sized rows), behind a size-dispatching topk_op;
ops/vanilla.py stays flashinfer-free via a topk_fn injection point.
CBA step fused with torch.compile (fullgraph + mark_dynamic, following
the Fusions pattern) and group-batched finalize D2H (one copy per
tensor 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

Top-beam parity with TRTLLMSampler across the (length_penalty,
diversity_rate, early_stopping, stop/VBWS) e2e combinations. Under
early_stopping 1 and 0, 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.
Sampler microbench vs TRTLLMSampler: ~45% faster on all comparable
configurations; the new features add no measurable overhead when enabled,
zero overhead when disabled. NB: measured before early_stopping=1 moved
onto the CBA path, so that mode's numbers predate the change and want a
re-run.
Combinations rejected at admission: disaggregated serving with
best_of > 1; a decreasing beam_width_array (only non-decreasing
schedules are defined, and finalize reads the array maximum); a best_of
that differs from max_beam_width.
Known 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-compatible or api-breaking. For api-breaking, include BREAKING in 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.

coderabbitai · 2026-07-22T11:01:58Z

Note

Reviews paused

It 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 reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

@coderabbitai resume to resume automatic reviews.
@coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

▶️ Resume reviews
🔍 Trigger review

Walkthrough

Beam-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.

Changes

Beam Search Sampling

Layer / File(s)	Summary
Beam strategy and top-k wiring `tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py`, `tensorrt_llm/_torch/pyexecutor/sampler/sampler.py`, `tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py`, `tensorrt_llm/_torch/pyexecutor/llm_request.py`, `tensorrt_llm/sampling_params.py`, `docs/source/features/sampling.md`	Beam-search parameters are resolved, grouped, documented, and passed to selected kernels; top-k dispatch and beam-width indexing were added.
Candidate beam selection `tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py`, `tests/unittest/_torch/sampler/test_beam_search.py`	Two-stage candidate top-k applies length and diversity adjustments, tracks predecessor beams, reorders stop-window state, and preserves raw cumulative log probabilities.
CBA execution and beam history `tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py`, `tensorrt_llm/_torch/pyexecutor/sampler/sampler.py`, `tests/unittest/_torch/sampler/test_beam_search.py`	CBA buffers, compiled step computation, EOS and stop-word handling, slot replacement, completion detection, and CBA-based history finalization were added.
Beam-search behavior validation `tests/unittest/_torch/sampler/test_beam_search.py`	Tests cover metadata updates, ranking adjustments, top-k parity, CBA transitions, stop-word harvesting, and stop-window reordering.

Estimated code review effort: 4 (Complex) | ~60 minutes

Possibly related PRs

NVIDIA/TensorRT-LLM#15542: Touches the same PyTorch sampling strategy and grouped execution pathways.

Suggested labels: api-compatible

Suggested reviewers: qijune, cascade812, reasonsolo, arysef, lfr-0531

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)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly summarizes the main TorchSampler beam-search enhancements and uses the required ticket and feature format.
Description check	✅ Passed	The description explains the problem, implementation, behavior changes, optimizations, results, and relevant tests; coverage is documented despite no separate Test Coverage heading.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

ixlmar

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.

coderabbitai

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

📥 Commits

Reviewing files that changed from the base of the PR and between db133d0 and 5e463e4.

📒 Files selected for processing (4)

tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py
tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
tensorrt_llm/_torch/pyexecutor/sampler/sampling_utils.py
tests/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

…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 · 2026-08-10T14:54:06Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-10T15:00:50Z

PR_Github #65050 [ run ] triggered by Bot. Commit: 5e395c6 Link to invocation

…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 · 2026-08-10T15:14:03Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-10T15:19:36Z

PR_Github #65058 [ run ] triggered by Bot. Commit: 238c420 Link to invocation

tensorrt-cicd · 2026-08-10T15:23:50Z

PR_Github #65050 [ run ] completed with state ABORTED. Commit: 5e395c6

Link to invocation

tensorrt-cicd · 2026-08-10T20:01:59Z

PR_Github #65058 [ run ] completed with state SUCCESS. Commit: 238c420
/LLM/main/L0_MergeRequest_PR pipeline #52865 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

zhaoyangwang-nvidia · 2026-08-11T01:30:23Z

/bot run

tensorrt-cicd · 2026-08-11T01:36:44Z

PR_Github #65162 [ run ] triggered by Bot. Commit: 238c420 Link to invocation

tensorrt-cicd · 2026-08-11T04:59:30Z

PR_Github #65162 [ run ] completed with state FAILURE. Commit: 238c420
/LLM/main/L0_MergeRequest_PR pipeline #52956 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

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>

zhaoyangwang-nvidia · 2026-08-11T06:40:03Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-11T06:45:33Z

PR_Github #65222 [ run ] triggered by Bot. Commit: 5f4906a Link to invocation

zhaoyangwang-nvidia · 2026-08-11T09:20:47Z

Hi @QiJune @hyukn could you help to take a look of this PR?

tensorrt-cicd · 2026-08-11T15:31:22Z

PR_Github #65222 [ run ] completed with state FAILURE. Commit: 5f4906a
/LLM/main/L0_MergeRequest_PR pipeline #53007 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

zhaoyangwang-nvidia · 2026-08-12T00:00:26Z

/bot run

tensorrt-cicd · 2026-08-12T00:06:17Z

PR_Github #65426 [ run ] triggered by Bot. Commit: 5f4906a Link to invocation

QiJune

LGTM

hyukn

LGTM for perf part. Thanks.

github-actions Bot assigned zhaoyangwang-nvidia Jul 20, 2026

zhaoyangwang-nvidia force-pushed the beam-length-penalty branch 3 times, most recently from a0cc88a to bcc284f Compare July 21, 2026 06:54

zhaoyangwang-nvidia changed the title ~~[None][feat] Support beam search length_penalty and diversity_rate in TorchSampler~~ [13234][feat] Support beam search length_penalty/diversity_rate and validate early_stopping in TorchSampler Jul 21, 2026

zhaoyangwang-nvidia force-pushed the beam-length-penalty branch 2 times, most recently from d246f47 to f11c23f Compare July 21, 2026 09:36

zhaoyangwang-nvidia changed the title ~~[13234][feat] Support beam search length_penalty/diversity_rate and validate early_stopping in TorchSampler~~ [TRTLLM-13234][feat] Support beam search length_penalty/diversity_rate and validate early_stopping in TorchSampler Jul 21, 2026

zhaoyangwang-nvidia changed the title ~~[TRTLLM-13234][feat] Support beam search length_penalty/diversity_rate and validate early_stopping in TorchSampler~~ [TRTLLM-13234][feat] Support beam search length_penalty/diversity_rate and early_stopping in TorchSampler Jul 21, 2026

zhaoyangwang-nvidia force-pushed the beam-length-penalty branch 7 times, most recently from d5dab0e to 7167148 Compare July 22, 2026 10:51

zhaoyangwang-nvidia changed the title ~~[TRTLLM-13234][feat] Support beam search length_penalty/diversity_rate and early_stopping in TorchSampler~~ [TRTLLM-13234][feat] Complete TorchSampler beam search: length_penalty, diversity_rate, early_stopping, VBWS, and CBA performance Jul 22, 2026

zhaoyangwang-nvidia marked this pull request as ready for review July 22, 2026 10:54

zhaoyangwang-nvidia requested review from a team as code owners July 22, 2026 10:54

zhaoyangwang-nvidia requested review from JunyiXu-nv, YihuiLu512, arysef and laikhtewari July 22, 2026 10:54

QiJune requested a review from lori-ren July 23, 2026 01:59

ixlmar reviewed Jul 23, 2026

View reviewed changes

zhaoyangwang-nvidia force-pushed the beam-length-penalty branch from 7167148 to 4721a74 Compare July 24, 2026 14:17

coderabbitai Bot reviewed Jul 27, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/pyexecutor/sampler/sampler.py Outdated

Comment thread tensorrt_llm/_torch/pyexecutor/sampler/sampler.py Outdated

zhaoyangwang-nvidia requested a review from a team as a code owner July 27, 2026 09:03

zhaoyangwang-nvidia force-pushed the beam-length-penalty branch from 66afaaf to 5e395c6 Compare August 10, 2026 14:54

zhaoyangwang-nvidia added 5 commits August 10, 2026 08:13

zhaoyangwang-nvidia force-pushed the beam-length-penalty branch from 5e395c6 to 238c420 Compare August 10, 2026 15:13

Shixiaowei02 approved these changes Aug 11, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/pyexecutor/py_executor.py

Comment thread tests/unittest/_torch/sampler/test_beam_search.py

Shixiaowei02 added the ci: full pre-merge approved label Aug 11, 2026

zhaoyangwang-nvidia mentioned this pull request Aug 11, 2026

[TRTLLM-11628][perf] Batch the beam-search finish-reason reduction #17494

Draft

4 tasks

QiJune approved these changes Aug 12, 2026

View reviewed changes

hyukn approved these changes Aug 12, 2026

View reviewed changes

zhaoyangwang-nvidia enabled auto-merge (squash) August 12, 2026 02:39

Uh oh!

Conversation

zhaoyangwang-nvidia commented Jul 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Dev Engineer Review

QA Engineer Review

Description

New support

Behavior changes

Optimizations

Results

PR Checklist

GitHub Bot Help

Uh oh!

coderabbitai Bot commented Jul 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Reviews paused

Walkthrough

Changes

Sequence Diagram(s)

Uh oh!

ixlmar left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

zhaoyangwang-nvidia commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

zhaoyangwang-nvidia commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

tensorrt-cicd commented Aug 10, 2026

Uh oh!

zhaoyangwang-nvidia commented Aug 11, 2026

Uh oh!

tensorrt-cicd commented Aug 11, 2026

Uh oh!

Uh oh!

Uh oh!

tensorrt-cicd commented Aug 11, 2026

Uh oh!

zhaoyangwang-nvidia commented Aug 11, 2026

Uh oh!

tensorrt-cicd commented Aug 11, 2026

Uh oh!

zhaoyangwang-nvidia commented Aug 11, 2026

Uh oh!

tensorrt-cicd commented Aug 11, 2026

Uh oh!

zhaoyangwang-nvidia commented Aug 12, 2026

Uh oh!

tensorrt-cicd commented Aug 12, 2026

Uh oh!

QiJune left a comment

Choose a reason for hiding this comment

Uh oh!

hyukn left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

zhaoyangwang-nvidia commented Jul 20, 2026 •

edited

Loading

coderabbitai Bot commented Jul 22, 2026 •

edited

Loading