[TRTLLM-11628][perf] Batch the beam-search finish-reason reduction - #17494
Draft
zhaoyangwang-nvidia wants to merge 37 commits into
Draft
[TRTLLM-11628][perf] Batch the beam-search finish-reason reduction#17494zhaoyangwang-nvidia wants to merge 37 commits into
zhaoyangwang-nvidia wants to merge 37 commits into
Conversation
…y, diversity_rate, early_stopping, VBWS, and CBA performance Implement length_penalty and beam_search_diversity_rate for the PyTorch sampler's beam search, matching the C++ decoder semantics, and add the exhaustive early_stopping modes backed by a candidate-beams array (CBA). Beam-search code moves into its own sampler.beam_search module behind a BeamSearchHandler, mirroring how token_ban / top_p_decay / finish_reasons are organized. - length_penalty: candidates ranked by (cum_log_prob + diversity_rate * source_beam_index) / gen_len**penalty via a two-stage top-k that never touches the vocab axis. Stored cum_log_probs stay raw. - early_stopping: TRUE (default) stops once beam_width finished candidates exist; FALSE and NEVER keep a pool of finished candidates and differ only in the bound used for a beam's best attainable score. - Variable beam width: per-iteration widths are honored, and get_beam_width_by_iter overrides the C++ binding, which reads past the end of the user array once decoding outruns it. - CBA performance: the step is fused and its finalize D2H batched. The CBA tensors are allocated on first use rather than for every beam-enabled engine, and the beam-search buffers stay allocated only while they are needed. Serving defaults are left unset so requests over HTTP keep the engine's beam-search defaults, as they did before this change: the OpenAI-compatible server used to send length_penalty=1.0 and early_stopping=False unconditionally, which was harmless only because the Torch sampler ignored both. Two combinations are rejected at admission rather than silently mishandled: a beam width below max_beam_width (the attention metadata is stamped with max_beam_width while the generation rows are laid out at the per-request width, and the scheduler cannot keep widths from mixing within a batch), and disaggregated serving with an exhaustive early_stopping mode (the finished-candidate pool the context server can populate is not part of the handoff). Both are tracked in TRTLLM-14792. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…test test_smaller_beam_width asserted on ".*exceeds max_beam_width.*", but _validate_request only ever raises "... is not equal to max_beam_width ...". The assertion could not match, so both parametrizations failed once the test was actually executed. The test class carries @force_ampere, so these cases skip on Hopper CI and the mismatch went unnoticed. Verified on H200 with TLLM_TEST_IGNORE_ARCH=1: the two cases now pass, and they still fail when the rejection in _validate_request is disabled. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…ma, VBWS tests _validate_request tested the bound method is_generation_only_request instead of calling it. is_context_only_request is a property but is_generation_only_request is a plain method, so the bound method was always truthy and the disaggregated-serving check matched every request: exhaustive early_stopping was rejected in regular serving too. The existing test only asserted that disagg rejects the combination, which an always-true condition also satisfies; add the missing direction and verify it fails before the fix. Narrow the HTTP early_stopping schema to bool | "never" | None, mirroring the HuggingFace interface, and translate to the engine's integer encoding in the protocol layer. Values such as 100 were previously accepted and silently treated as NEVER. Integer 1/0 still validate as True/False, so existing clients are unaffected. Add Variable-Beam-Width-Search coverage, which had none: beam_width_array was never driven through the engine. Unit tests cover widening, narrowing, holding the last width once decoding outruns the array, and equivalence between a constant array and a fixed width. One test pins the divergence from the C++ formula, which clamps with the global kMaxBeamWidthArrayLength and reads past the end of the user array (observed: 0, 32, 849); it fails if C++ is fixed, so the Python override is not dropped as redundant. A kernel test covers a width transition with length_penalty, where per-beam lengths must follow the beam permutation. An end-to-end test drives the full scheduler -> ModelEngine -> TorchSampler path. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…e2e test The VBWS end-to-end test only checked the returned beams, which a request running at a constant width the whole time would satisfy just as well as one that actually walked beam_width_array. Record the width the engine hands out at each decoding iteration by wrapping LlmRequest.get_beam_width_by_iter -- the accessor the scheduler, ModelEngine and sampler all go through -- and assert it follows beam_width_array and then holds at the last entry once decoding outruns the array, with every width in the array exercised. Verified against a real model (TinyLlama-1.1B, beam_width_array=[2,3,4], max_tokens=5), where the engine reports widths 2, 3, 4, 4 over decoding iterations 1-4. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…th array length getBeamWidthByIter clamped the decoding-iteration index with the global kMaxBeamWidthArrayLength constant instead of the length of the array the user actually supplied. Once decoding ran longer than that array, the index walked past its end and the returned beam width was garbage read out of bounds. The C++ micro-batch scheduler calls this method for every generation request. With a garbage width it never admits the request into the generation batch again, so the request stays GENERATION_IN_PROGRESS forever, decoding_iter stops advancing and generate() never returns -- the executor loop spins while the request is silently starved. Clamp with the array's own size instead, so decoding past the end holds the last width, matching the Python LlmRequest override. Also guard the empty-array case, which was previously an unchecked index. Reproduced and verified on H200 with a beam_width_array whose length is shorter than max_tokens: [2,3,4]/max_tokens=5, [2,3,4,4]/max_tokens=6 and [4,4,4]/max_tokens=5 all hung before and complete after, while [2,3,4]/max_tokens=3 and [2,3,4,4]/max_tokens=5 (which never outrun the array) passed both before and after. Note the hang is not specific to varying widths -- a constant [4,4,4] array hangs just as well. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…VBWS e2e test max_tokens was capped at 5 to work around a hang, leaving only two decoding iterations on the clamp that holds the last width once decoding outruns beam_width_array. The workaround did not actually avoid the hang either, since the array is three entries long and anything above three triggered it. That hang is fixed (getBeamWidthByIter now clamps with the array's own length), so decode to 12 tokens instead: nine of the twelve iterations now exercise the clamp, which is the part that used to read out of bounds. Replace the stale "not yet diagnosed" note with the actual cause and record that the test needs the C++ fix -- against an older libtensorrt_llm.so it hangs rather than fails. Verified on H200: max_tokens 8 and 12 both complete, and 12 completed on four consecutive runs, each returning four distinct beams of exactly twelve tokens. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
… width check The mixed-beam-width guard filtered only CUDA-graph dummies, but dummy requests come in three kinds and the other two also carry a width of their own: attention-DP and warmup dummies are built at width one. With attention-DP enabled, such a dummy joining a beam-search generation batch reports width one against the real requests' width and aborts the whole batch. Filter on is_dummy, which covers all three flags. Also update test_vbws_cpp_formula_reads_past_array_end, which asserted that the C++ clamp still diverges from the Python one past the end of beam_width_array. That divergence was the hang fixed earlier in this branch, so the test now pins the opposite: the two must agree. Renamed accordingly. It fails against a libtensorrt_llm.so built before that fix (C++ returns 0 where Python returns the last width), which is the intended signal. Reported by Shixiaowei02 in review. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
ModelEngine lays out generation rows at the static admission width (py_beam_width), but the sampler located a request's logits by accumulating per-iteration widths. Those agree for a fixed beam width and diverge under a variable beam width array: at an iteration narrower than max_beam_width, every request after the first read another request's rows. logits.view() accepts any shape whose element count divides, so the result was silently wrong rather than an error. Neither existing guard catches this. Admission compares beam_width, which checkBeamWidthArray raises to the array maximum, so a VBWS request looks exactly like a fixed max-width one. The scheduler and the ModelEngine check both compare per-iteration widths across the batch, which are equal when requests advance in lockstep -- the very case that misreads. Offset by py_beam_width instead, matching the layout, and slice down to the per-iteration width in the beam-search ops where the live beams are consumed. Thread that stride through as row_stride, defaulting to beam_width_in so non-VBWS paths are unchanged. TRTLLMSampler already offsets by the static width, which is why it is unaffected. NB: row_stride is appended last in the BeamSearch tuple because _common_fields() reads the preceding fields by position. Reported by Shixiaowei02 in review. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
Adding row_stride to BeamSearchStep.__init__ left the abstract-class test constructing it with the old five positional arguments, which CI's mypy flagged as a missing argument. The test only pins that the base class cannot be instantiated, so supply the extra width and keep the contract. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
early_stopping=True was served by a separate path that freezes a beam in its slot once it finishes, and stopped when every slot held a finished beam. A comment claimed that was equivalent to the C++ decoder. It is not: C++ counts hypotheses accumulated in the candidate-beams array, and because a finished candidate vacates its slot there, one lineage can contribute several hypotheses. Freezing instead pins the slot, so that lineage stops being explored and the run needs beam_width distinct slots to finish. vLLM does what C++ does -- 2*beam_width candidates per step, finished beams moved to a completed pool, the slot refilled -- so the frozen-slot behaviour was ours alone, on the default mode. Route every mode through beam_search_sampling_batch_cba. TRUE differs only in the done verdict: stop as soon as the pool is full, without weighing what is still attainable, matching beamStage3Kernel. Beam search under disaggregated serving is rejected as a consequence. The pool is not part of the handoff -- ContextPhaseParams carries only the first generated tokens -- and is cleared on the generation side, so a completion the context phase found would be silently dropped. This now applies to every mode rather than the exhaustive ones only, and to both samplers: the C++ decoder assigns numBeamsCBA with no mode check, so TRTLLMSampler was always affected. Rejecting it costs the beam score handoff added earlier in this branch; the corresponding e2e test now pins the rejection instead. Also reject a decreasing beam_width_array. The documented VBWS semantics only cover widening, and both samplers depend on it: a step writes the leading beam_width_out rows of the beam state while finalize reads py_beam_width (the array maximum) of them, so a narrowing array returns beams whose ancestry and cumulative log-probs are left over from an earlier, wider step. The C++ finalize path has the same gap. No other engine implements variable beam widths at all -- vLLM and HF take a scalar, SGLang has no beam search -- so there is no reference semantics to follow for narrowing. NB: not verified end to end. The prebuilt libraries in my environment no longer match the current bindings, so every engine-level test fails at KVCacheManager construction regardless of this change. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
Two comments in TorchSampler still described the removed frontier: one claimed early_stopping == TRUE runs on the frozen-slot path, the other that the default mode never needs the candidate-beams-array tensors. Both now do. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…le arity strategy_grouping_key() matched the BeamSearch tuple by exact arity, so appending row_stride to it sent every beam-search strategy to the "Unsupported strategy encountered" branch: TorchSampler failed every beam search request. Match the leading fields and ignore the rest, and pin it with a test -- the existing unit tests build the ops and step classes directly, so none of them exercise the grouping layer, which is why this went unnoticed until a real engine ran. Fold the early_stopping == TRUE done verdict into the same expression as the other modes instead of branching. `early_stopping` is a plain int in _cba_step_math, so a Python branch on it is another Dynamo guard on a fullgraph-compiled function. Split test_create_beam_history, which drove _prepare_beam_history end to end and no longer describes its output: finalization now always takes the CBA path, which merges the finished-candidate pool into the active beams and reorders by normalized score, while the test computed a per-beam expectation from cache_indirection alone. It becomes test_gather_beam_path_follows_cache_indirection, covering the ancestry gather both paths share, plus test_cba_finalize_merges_pool_and_orders_by_score for the merge and ordering. The latter drives the real _prepare_beam_history_cba rather than restating its arithmetic. Verified on H200 against a rebuilt libtensorrt_llm.so: 43 passed, 0 failed. TorchSampler and TRTLLMSampler now return three of four identical beams for early_stopping True and False; the fourth differs by the EOS convention (TorchSampler keeps a trailing EOS, so a short finished hypothesis ranks differently), which predates this change. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…istories The sampling op pads its output row out to the store's beam width, but update_requests() appended a token for every column of that row. On a variable-beam-width step the trailing columns hold BEAM_SEARCH_PAD_TOKEN, so the sentinel landed in the request's token history. Finalization later rewrites every beam from the corrected paths, which is why the final outputs look right and no existing test noticed -- but the padded history is visible to streaming consumers and to anything reading get_tokens() mid-flight, such as the token-ban suffix matching. Append only the beams the step produced. Fix two more sites that read a beam width where the row layout uses the static one, both consequences of moving the logits offsets to the static stride: - the beam-search temperature was repeat_interleave'd by beam_width_in while the op receives row_stride rows, so a widening array failed outright with a shape mismatch - _execute_logit_post_processors advanced logits_row_offset by the per-iteration width, so with several generation requests in a batch every request after the first rewrote another request's logits rows Extend test_beam_search_vbws_e2e to inspect the requests after every update_requests() and assert the sentinel is absent. Verified it fails against the previous implementation (beam_width_array=[2,3,4], max_beam_width=4 leaked the sentinel ten times) and passes after, with the final outputs unchanged. Reported by QiJune in review. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
Routing every early_stopping mode through the candidate-beams-array path left a dozen comments describing the frontier it replaced. The worst directly contradicted the fix in the preceding commit: _pad_next_tokens claimed the padding "is never consumed", which is exactly what update_requests() was doing. The rest still gated the CBA state on early_stopping != TRUE, or named beam_search_sampling_batch as the op in use. Delete RegularBeamSearchStep, added by this branch and dispatched to by nothing since the unification. Delete two tests along with it: test_beam_search_sampling_batch_disagg_handoff covers a path admission now rejects outright, and test_beam_search_sampling_batch_reorders_stop_window duplicates test_beam_search_cba_reorders_stop_window on the live op. Annotate what is kept. beam_search_sampling_batch, _check_beam_search_stop_criteria and the pool-free branch of _prepare_beam_history are unreachable but predate this branch, as does _update_sampler_state_for_disagg_gen_request, which is the seeding half of a disaggregated handoff that becomes live again once the finished- candidate pool is transferred. None of them said so; they do now. Still to do, deliberately left out: removing beam_search_sampling_batch itself, which needs its four remaining unit tests (basic, length_penalty, diversity_rate, VBWS width transition) ported onto the CBA op. Their expected values have to be re-derived, since the two ops differ in how a finished beam is treated, and I would rather do that against a running engine than by reading the implementation. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
The HTTP schema change is the one users can hit without touching their
request: length_penalty defaulted to 1.0 and now defers to the engine
default of 0.0, so a beam-search request that never set it stops
normalizing scores by length. sampling.md now says so and tells readers
how to keep the old ranking, and covers the new false/true/"never"
spelling.
sampling.md also described early_stopping as "0, and other values for
intermediate heuristics". There are no intermediate heuristics:
from_raw maps everything outside {0, 1} to "never". Replaced with the
actual three states, and mirrored the same wording plus the
length_penalty formula and the beam_width_array constraint into the
SamplingParams docstrings, which had none of it.
Document the three combinations that now raise at admission --
disaggregated serving, a decreasing beam_width_array, and a best_of
that differs from max_beam_width -- none of which appeared anywhere in
docs/.
Fix comments that no longer match the code:
- LlmRequest.get_beam_width_by_iter still justified itself by a C++ bug
this branch fixed. The override is still wanted, but for a different
reason: the binding is not virtual, so Python callers would otherwise
pick up whatever prebuilt library is present.
- _common_fields claimed grouping guarantees a single row_stride per
group. It does not -- row_stride is deliberately excluded from the
grouping key. It holds because admission pins every request to
max_beam_width, which is what the comment now says.
- BeamSearchStore said the default early_stopping never touches the CBA
tensors, contradicting ensure_cba three screens below.
- A four-line comment appeared twice in a row.
The assert_no_cuda_sync removal was justified by softmax and a
scalar-float division synchronizing. Measured: they do not. The guard
trips on the caching allocator's first allocation for a shape, and all
three ops pass once the allocator is warm -- which is what
run_test_with_warmup exists for. Comment corrected.
Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
test_beam_search_vbws_e2e drove one schedule with default parameters, but every width-related bug fixed in this branch sat where a variable beam width array crosses another feature: per-beam generated lengths under length_penalty, source-beam ranks under diversity_rate, and the candidate pool width under the exhaustive early_stopping modes. Parametrize it over those combinations, plus a constant [4, 4, 4] array as a control that separates "the VBWS plumbing is broken" from "changing width is broken" -- which is how the decoding hang was localized. Raise the Dynamo recompile limit for the duration of the test. Each parametrization builds an engine in the same process and the CBA step is compiled with fullgraph=True, so the per-code-object recompile count is exhausted by the last case and compilation fails hard instead of falling back. Verified: without this, the run fails at [es_never] while that case passes on its own; with it, all six pass in one process. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…y path The existing length_penalty test drives beam_search_sampling_batch, where a finished beam stays in its slot and competes with the live ones, so the penalty is observed by watching the two swap. That is not what the penalty does on the CBA path: candidate ranking there does not use it at all, and a finished beam leaves the slots for the pool. Porting the test across would have meant re-deriving its expectations from the CBA implementation, which is not a check on anything. Add a test built around what the penalty does there instead. Two beams finish on the same step with different frozen lengths -- the shorter one ahead on raw cumulative log-prob, the longer one on per-token score -- and the assertions are on the pool: raw cum_log_probs are identical with and without the penalty, the normalized scores are not, and each equals its entry's cum divided by that entry's recorded length. Verified the assertions bite: zeroing the exponent at the pool insertion fails both parametrizations. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
The docstring said to keep the function free of data-dependent shapes and in-place mutation, but the rest of what fullgraph=True demands lived only at the call site: dim 0 of six inputs is mark_dynamic and must never be read as a size or branched on, snap_arange is only maybe_mark_dynamic so specializing on it is allowed, and the plain-int arguments are static and cost a compilation per distinct combination. That is the kind of contract a torch upgrade breaks, and the next person to edit this function would not see it. Move it into the docstring. Also note that _beam_step_preprocess is eager only -- it writes the cache indirection buffer in place -- so it does not get pulled into the compiled region by someone looking for more fusion. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
Nothing has dispatched to beam_search_sampling_batch since every early_stopping mode moved onto the candidate-beams-array op. Keeping it was justified while its unit tests still covered ranking behaviour the CBA path had no equivalent for; that gap is now closed -- test_beam_search_cba_length_penalty_orders_pool covers the penalty where it actually applies, and the VBWS width transitions it exercised are covered end to end by the parametrized engine test. Remove the op, _check_beam_search_stop_criteria (its stop criterion, with no other caller), the unreachable branch of _prepare_beam_history, and _request_uses_cba, which had become a constant-true predicate that implied a pool-free path still existed. Their three unit tests go too: they assert that a finished beam stays in its slot and emits a pad token, which is the behaviour this branch deliberately replaced. _pad_next_tokens stays -- the CBA op uses it as well. Verified on H200: 38 unit tests and all six VBWS engine cases pass. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
Deleting beam_search_sampling_batch left two store fields with writers
but no readers. beam_gen_lengths held the per-beam generated length that
op froze when a beam finished; the CBA step derives lengths from
seq_lens - prompt_lens instead, so the field was allocated, reset every
step, seeded by the disagg handoff and read by nobody. seq_offsets
existed to flatten (batch, beam) pairs for that op's indexing; the CBA
step uses advanced indexing and never flattens.
_postprocess_beam_logprobs and the two dataclasses it consumed
(_BeamHistoryTensors, _BeamHistoryLogProbsSlices) belonged to the
finalize branch removed alongside the op and have no callers.
Also fix comments the same removal invalidated, including a CBAState
docstring left half-overwritten by an earlier edit ("present only for
requests using / every beam-search request"), a claim that finished
beams stay frozen in their slots "rather than in a separate pool" (the
CBA is that pool), and two references to the deleted op.
One test change is a correctness fix rather than fallout: the CBA
length_penalty test set beam_gen_lengths and commented that it was
establishing a shorter frozen length for one beam. Nothing reads that
field, so the line was decorative -- the lengths that drive the assertion
come from seq_lens - prompt_lens.
Kept: beam_candidate_topk's length_penalty/cand_gen_lengths parameters,
which no production caller passes but test_beam_candidate_topk_equivalence
uses to check the two-stage top-k against a naive full-matrix adjustment;
and predecessor_beams, which has no production reader but is asserted on
by test_logits_logprobs.
Verified on H200: 38 unit tests and all six VBWS engine cases pass.
Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
The speculative beam-history D2H path predicts whether a step is likely terminal and, on a hit, snapshots the beam history off the device. The snapshot is a single batched copy covering the whole group, so the skip decision cannot be made per request: skipping for some requests while copying for others would drop the histories of the skipped ones. Gate the copy on the group instead. A step is skipped only when every request is predicted non-terminal; a single possible finisher makes the group issue the copy it would have cost anyway. The fallback for a mispredicted step takes its snapshot synchronously rather than consulting the host-side mirror of finish reasons, which lags one step behind and would misread a stop-word hit as unfinished. Add a test pinning the predictor to disagree across requests, asserting that the mixed verdict lands on "copy" and that outputs still match the feature-off run. The gating short-circuits on the feature flag, so the predictor is never consulted and the new path is unreachable when enable_speculative_beam_history_d2h is False (the default). Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
Test and contract fixes from review of the beam-search work: - The VBWS admission test asserted a local re-implementation of the non-decreasing check instead of calling _validate_request, so deleting the production check would not have failed it. Drive the real validator. - beam search reached flashinfer's radix top-k unconditionally on real vocabularies, making the default beam path require flashinfer. Guard on IS_FLASHINFER_AVAILABLE and fall back to torch.topk. - Drop the comment describing a generated-length counter that no longer exists; length_penalty now normalizes by seq_lens - prompt_lens, which counts the handed-off token because it occupies the first generated slot. - _prepare_beam_history had degenerated to a two-line forward whose two parameters were unused, while the call site evaluated a per-request GPU index op only to discard it. Call _prepare_beam_history_cba directly. - get_beam_width_by_iter raised IndexError on an empty beam_width_array where the C++ side guards; fall through to the base implementation. - Reject a negative length_penalty at admission: a negative exponent inverts the beam ordering. - Two docs claimed more than the code does: admission does not catch mixed per-iteration VBWS widths, and the C++ past-the-end read the test describes in the present tense is fixed in this PR. - The width-walking test pinned narrowing behaviour that admission rejects, contradicting the rejection test in the same file. - Restore the @_kernel_test decorator on two CBA tests that silently ran eager on CPU, and drop a duplicated one. - Fill CBA test buffers with a poison value rather than zeros (zero is a plausible token id, length or beam index) and assert that rows outside seq_slots come back bit-identical. - Document that the CBA step keeps candidate scores diversity-free, matching HF (diversity is a logits processor applied before log_softmax) and vLLM (ranks purely by length-normalized cum_logprob). The C++ kernels fold the diversity term into the CBA insertion score and the done verdict, diverging from all three. - Document the execution contract of beam_search_sampling_batch_cba: in-place mutation, seq_slots-only writes, mark_dynamic on caller tensors, and the compiled/eager split. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com> - Raise Dynamo's recompile limit for the whole module instead of a single test. The cap is counted per code object across the process, so the compiled CBA step exhausted it partway through the file and every later case failed to compile under fullgraph; restoring @_kernel_test on the two eager tests would otherwise have pushed them into that same wall. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…BA path `args.finished_beams` is the finish handler's `first_finish_reasons`: it records the reason each beam *first* finished by, and the finish handler only ever fills entries that are still NOT_FINISHED. The CBA step wrote the done verdict over the whole row unconditionally, so on every step where the request is not yet done it cleared that record back to zero. On this path a beam that hits a stop word does not freeze -- it vacates its slot and the request keeps generating -- so its STOP_WORDS entry is the only surviving evidence of why it ended. Clearing it made the request report the reason it eventually stopped for, LENGTH, and a beam-search request with stop_token_ids returned finish_reason "length" instead of "stop". Write the verdict only where it is done, leaving earlier reasons intact. Found by running the beam-search e2e matrix one case per process: 28 stop_token_ids cases failed on `assert 'length' == 'stop'` and passed on the merge-base, which located the regression in this PR rather than in the recompile-limit noise the batched run was producing. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
The previous commit stopped the CBA step from clearing first_finish_reasons but still flooded the row with END_ID once the request was done, so every beam reported the verdict rather than its own reason. Fill only entries still at NOT_FINISHED, and use LENGTH for them: those beams end because the pool can no longer be beaten, not because they hit a token, which is what the pool-free path reported for the same situation. A stop_token_ids beam search now returns [STOP_WORDS, LENGTH] for its two beams, matching the merge-base. Four checks in the e2e test assumed a beam that never hit a stop token runs to max_tokens. Under early_stopping=1 -- the default, HF's `True`, which stops "as soon as there are num_beams complete candidates" -- the request can stop while a beam is shorter, so indexing cache_indirection at max_tokens - 1 ran off the end of the buffer and the logits/logprobs/token counts came up short. Compare against the length the beam actually produced; a diverging beam still fails on the prefix, only the unreached tail is dropped. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
- Drop the leading underscore from finalize_beam and prepare_beam_search. Both have consumers outside beam_search.py (the sampler, the beam-search tests, and the host profiler's function-name list), so the prefix claimed a privacy the names did not have. Same point lori-ren raised on sampler_strategy.py. - Rename the `olp` alias to beam_log_probs and say why the variable exists at all: advanced indexing returns a copy, so the scatter has to be written back. - Point the CBA-allocation comment at BeamSearchStore.ensure_cba and its caller; BeamSearchHandler.ensure_cba_for_requests never existed. - Delete BeamSearchHandler._log_probs_store and ._new_tokens, assigned in the constructor and never read since the pool-free path went away, along with the two constructor parameters and the arguments the sampler passed. - Record why BeamSearchEarlyStop.from_raw stays more permissive than the OpenAI-compatible server, which rejects values outside its tri-state: sampling_config.early_stopping is a plain int that also arrives from the C++ runtime, so folding unknown values into the nearest mode keeps a path working that already accepted them. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
The row_stride query added to sampler_common reads request.py_beam_width, which TestTopPDecay's SimpleNamespace stub does not define, so validate_request raised AttributeError before reaching the ValueError the test asserts on. These cases are single-beam; set the width to 1. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…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>
…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>
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>
_handle_finish_reasons_impl reduced one request's row of first_finish_reasons at a time, so every beam-search request in the batch paid a slice, a compare, a sum and a tensor-to-bool conversion to answer a question about at most max_beam_width integers. The dispatch overhead of those per-request ATen calls dominates the arithmetic they perform, and it lands on the host critical path of every decode step. Reduce the whole tensor once instead: count, per slot, how many leading beams have finished. The per-request check then reduces to comparing that count against the request's own beam width, which is a plain list lookup, and only the request updates -- the state assignment and the per-beam set_finished_reason calls, both of which cross into pybind -- keep looping. Counting the finished prefix rather than the finished beams keeps the result independent of what the columns past a request's beam width hold, so slots with differing beam widths share one reduction without relying on the padding columns being reset. 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
1 participant
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.
Description
_handle_finish_reasons_implreduced one request's row offirst_finish_reasonsat atime, so every beam-search request in the batch paid a slice, a compare, a sum and a
tensor-to-bool conversion to answer a question about at most
max_beam_widthintegers.The dispatch overhead of those per-request ATen calls dominates the arithmetic they
perform, and it lands on the host critical path of every decode step.
This reduces the whole tensor once instead: count, per slot, how many leading beams have
finished. The per-request check then reduces to comparing that count against the request's
own beam width, which is a plain list lookup, and only the request updates -- the state
assignment and the per-beam
set_finished_reasoncalls, both of which cross into pybind --keep looping.
Counting the finished prefix rather than the finished beams keeps the result independent
of what the columns past a request's beam width hold, so slots with differing beam widths
share one reduction without relying on the padding columns being reset.
Dependency
main. The diff shown here includes that PR'scommits. It should be reviewed/merged after #16620 lands, at which point this branch will
be rebased.
Test Coverage
Added to
TestFinishReasonsintests/unittest/_torch/sampler/test_torch_sampler.py:test_finished_beam_prefix_lengths_matches_per_request_reduction-- exhaustiveequivalence against the reduction this replaces, over all 4-column/4-reason rows and
every beam width.
test_finished_beam_prefix_lengths_ignores_columns_past_beam_width-- padding columnsmust not complete a request, nor mask an unfinished beam.
test_handle_first_finish_reasons_completes_only_fully_finished_requests-- mixed beamwidths (2/4/1) in one batch; asserts which requests complete and the exact per-beam
reasons recorded.
(exhaustive + randomized comparison of the old and new formulas) and
pre-commitpasses,but a GPU environment was still building at the time this draft was opened. Marking as
draft until
test_torch_sampler.py -k FinishReasons,test_beam_search.pyandtest_beam_search_speculative_d2h.pyhave actually run green.PR Checklist
[JIRA][type] summaryJIRA: https://jirasw.nvidia.com/browse/TRTLLM-11628
🤖 Generated with Claude Code