From 65d4801490c9bc7b9ede87f2ad86744ef30be6b7 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Thu, 30 Jul 2026 19:02:51 -0700 Subject: [PATCH 01/36] [TRTLLM-13234][feat] Complete TorchSampler beam search: length_penalty, 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 --- docs/source/features/sampling.md | 17 + tensorrt_llm/_torch/pyexecutor/llm_request.py | 24 + .../_torch/pyexecutor/model_engine.py | 43 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 43 +- .../_torch/pyexecutor/sampler/beam_search.py | 1706 +++++++++++++++++ .../pyexecutor/sampler/ops/flashinfer.py | 21 + .../_torch/pyexecutor/sampler/ops/vanilla.py | 132 -- .../_torch/pyexecutor/sampler/sampler.py | 572 +----- .../pyexecutor/sampler/sampler_common.py | 32 +- .../pyexecutor/sampler/sampler_strategy.py | 298 ++- tensorrt_llm/sampling_params.py | 2 +- tensorrt_llm/serve/openai_protocol.py | 16 +- .../_torch/sampler/test_beam_search.py | 891 ++++++++- .../test_beam_search_speculative_d2h.py | 9 +- .../references/trtllm_serve_api.yaml | 16 +- 15 files changed, 3086 insertions(+), 736 deletions(-) create mode 100644 tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py diff --git a/docs/source/features/sampling.md b/docs/source/features/sampling.md index 4f8eb101f450..d8818fb80770 100644 --- a/docs/source/features/sampling.md +++ b/docs/source/features/sampling.md @@ -274,6 +274,23 @@ Parameter Configuration: - `n`: Controls the number of output sequences returned (can be less than `best_of`) - If `best_of` is omitted, the number of beams processed defaults to `n` - `max_beam_width` in the `LLM` class must equal `best_of` in `SamplingParams` +- `length_penalty`: Controls how beams of different lengths are compared. Candidate beams are + ranked by `cum_log_prob / length**length_penalty`, where `length` is the number of generated + tokens. The default (`0.0`) ranks beams by their raw cumulative log-probability, which favors + shorter sequences; values above `0.0` favor longer sequences. The `cumulative_logprob` values + returned with the outputs remain unnormalized. +- `beam_search_diversity_rate`: Encourages beams to diverge from each other. During beam + expansion, `diversity_rate * source_beam_index` is added to each candidate's ranking score, + boosting candidates that expand from lower-ranked beams so that the selected beams do not all + descend from the single strongest beam. Here `source_beam_index` is the rank of the beam a + candidate expands from among the current step's input beams, ordered by their cumulative + log-probability (`0` for the strongest beam, `1` for the next, and so on). The default (`0.0`) + disables the adjustment. +- `early_stopping`: Controls when beam search stops. With the default (`1`), generation ends as + soon as `best_of` finished candidates exist. The exhaustive modes (`0`, and other values for + intermediate heuristics) keep a pool of finished candidates and continue searching while an + unfinished beam could still outscore the worst of them (`0` bounds attainability with the + current length; other values with the maximum length when `length_penalty > 0`). The following example demonstrates beam search with a beam width of 4, returning the top 3 sequences: diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index c68b83c2abde..174101559fd9 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -868,6 +868,30 @@ def __init__( else: self._py_embedding_bias_1d = self.embedding_bias + def get_beam_width_by_iter(self, for_next_iteration: bool = False) -> int: + """Beam width of the current (or next) decoding step. + + Overrides the C++ binding for Variable-Beam-Width-Search: the C++ + implementation clamps the decoding-iteration index with the global + kMaxBeamWidthArrayLength constant (it assumes a padded array) and + reads past the end of the raw user array once decoding runs longer + than the array — returning garbage widths. Same formula, clamped + with the actual array length. + + NB: the C++ method is not virtual and the binding has no trampoline, + so this override only covers callers on the Python side; C++ callers + still use the unfixed formula. + """ + beam_width_array = self.sampling_config.beam_width_array + if beam_width_array is not None: + if beam_width_array and isinstance(beam_width_array[0], + (list, tuple)): + beam_width_array = beam_width_array[0] + iteration = self.decoding_iter + (1 if for_next_iteration else 0) + index = max(min(iteration, len(beam_width_array)) - 1, 0) + return int(beam_width_array[index]) + return super().get_beam_width_by_iter(for_next_iteration) + def set_exclude_last_generation_logits( self, exclude_last_generation_logits: bool): self.py_result.set_exclude_last_generation_logits( diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index bd832625d042..dcc55f586efb 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -5227,8 +5227,49 @@ def append_cross_attention_state(request: LlmRequest, _has_any_multimodal_request = any(r.py_multimodal_data is not None for r in generation_requests) if _n_gen > 0: - # All generation requests have the same beam width + # The whole batch is laid out with request 0's beam width: every + # generation request contributes exactly this many rows to + # input_ids / position_ids / sequence_lengths and to the logits the + # model returns. The sampler, in turn, locates a request's logits by + # accumulating the *per-request* beam widths + # (TorchSampler._select_generated_logits -> + # calculate_request_offsets). Both agree only while every request in + # the batch has the same beam width. + # + # Mixing widths would desynchronize the two: the sampler would read + # a request's rows at the wrong offset, and `logits.view(batch, + # beam_width_in, vocab)` succeeds for any shape whose element count + # divides, so the result is silently wrong rather than an error. + # Supporting mixed widths needs the forward path to emit a fixed + # max_beam_width stride and the sampler offsets to match; until + # then, fail loudly. beam_width = generation_requests[0].py_beam_width + # Admission pins every request to max_beam_width, but a + # variable-beam-width request narrows or widens per iteration, so + # the widths can still diverge mid-batch. Compare the + # *per-iteration* width: py_beam_width is fixed at admission and + # would be identical across those requests. CUDA-graph padding + # requests are built at the engine width and appended after + # scheduling, so they are excluded -- they carry no user request + # and would otherwise trip this on an ordinary padded batch. + real_requests = [ + req for req in generation_requests + if not req.is_cuda_graph_dummy + ] + iter_widths = { + req.get_beam_width_by_iter() + for req in real_requests + } + if len(iter_widths) > 1: + # NB: this aborts the whole batch, not just the offending + # requests -- ModelEngine has no per-request failure channel, + # and by this point the batch is already scheduled. Scoping the + # failure needs the scheduler to group by beam width in the + # first place, so that no such batch is formed; TRTLLM-14792. + raise ValueError( + "Generation requests in one batch must all have the same " + f"beam width; got {sorted(iter_widths)}. Mixed beam widths " + "within a batch are not supported yet (TRTLLM-14792).") # Pre-extend constant-value lists to avoid per-request append # overhead (saves ~3 append calls per request). diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 62b7c6488ee9..3ba1777e5937 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -87,6 +87,8 @@ ResourceManagerType, request_context) from .sampler import (AsyncWorkerMixin, Sampler, SamplerEvent, SampleState, SampleStateTensors, TRTLLMSampler) +from .sampler.beam_search import BeamSearchEarlyStop +from .sampler.sampler_common import _unwrap_singleton from .scheduler import (RequestScheduler, ScheduledRequests, SerializableSchedulerOutput, WaitingQueue, create_waiting_queue) @@ -4974,11 +4976,41 @@ def _validate_request(self, request: LlmRequest): # Validate beam width sampling_config = request.sampling_config if sampling_config is not None: + # Requests must run at exactly max_beam_width. + # + # TorchSampler can sample a narrower request (buffers are allocated + # at max_beam_width and it slices to the per-request width), but the + # layers around it are not ready: the attention metadata is stamped + # with max_beam_width while the generation rows are laid out at the + # per-request width, and the scheduler is not beam-width aware, so a + # narrower request can be batched with a wider one and fail at + # forward time. Keep rejecting until those agree; TRTLLM-14792. if sampling_config.beam_width != self.max_beam_width: raise ValueError( f"Request beam width {sampling_config.beam_width} " - f"is not equal to max_beam_width {self.max_beam_width}. This is not supported!" - ) + f"is not equal to max_beam_width {self.max_beam_width}. " + "This is not supported!") + + # Exhaustive early_stopping keeps a pool of finished candidates, + # which the context server can already populate on its first step. + # That pool is not part of the disaggregated handoff and is cleared + # on the generation side, so a completion the context phase found + # would be silently dropped -- reachable under aggregation but not + # under disaggregation. Reject the combination until the handoff + # carries the pool; TRTLLM-14792. + if (request.is_context_only_request + or request.is_generation_only_request): + early_stopping = _unwrap_singleton( + sampling_config.early_stopping) + if (early_stopping is not None + and BeamSearchEarlyStop.from_raw(early_stopping) + is not BeamSearchEarlyStop.TRUE): + raise ValueError( + f"Beam search early_stopping={early_stopping} is not " + "supported with disaggregated serving: the finished-" + "candidate pool is not transferred between the context " + "and generation servers. Use the default " + "(early_stopping=True).") # Check token ID ranges self._validate_token_id_range(request) @@ -6430,6 +6462,13 @@ def fail_request(message: str) -> bool: device=cum_log_probs.device, dtype=cum_log_probs.dtype) cum_log_probs[seq_slot, :beam_width].copy_(values) + + # The handoff carries one token already produced upstream. The + # per-beam generated-length counter is reset to zero when the request + # is admitted here, so seed it to 1: length_penalty normalizes by this + # counter, and leaving it at zero would divide by one less than the + # true length for the whole request. + beam_search_store.beam_gen_lengths[seq_slot, :beam_width].fill_(1) return True @staticmethod diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py new file mode 100644 index 000000000000..061861220e1a --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py @@ -0,0 +1,1706 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""PyTorch-native beam-search sampling kernels. + +The candidate-selection, beam-expansion, and candidate-beams-array (CBA) +exhaustive-early-stopping logic for TorchSampler beam search, split out of +``ops.vanilla`` so the regular and CBA paths live in one place. Pure tensor +functions plus their metadata dataclasses; no dependency on the +``sampling_utils`` interface. +""" + +from contextlib import AbstractContextManager, nullcontext +from dataclasses import dataclass +from enum import IntEnum +from typing import TYPE_CHECKING, Callable, NamedTuple, Optional, TypeAlias, cast + +import torch + +from tensorrt_llm._utils import nvtx_range, prefer_pinned +from tensorrt_llm.bindings.executor import FinishReason + +from ..llm_request import LlmRequest, LlmRequestState +from .logprobs import LogProbsStore, convert_logprobs_tensor_to_list, get_logprobs_from_request +from .ops.flashinfer import radix_topk_op +from .ops.vanilla import StrategyMetadata +from .sampler_common import _get_beam_width_in, _unwrap_singleton, int_tensor + +if TYPE_CHECKING: + # Type-only: the async-D2H copier is a private detail of sampler.py and is + # injected as a bound method, so importing it at runtime would create a cycle. + from .sampler import _SideStreamCopier + +BEAM_SEARCH_PAD_TOKEN = -1 + + +class BeamSearchEarlyStop(IntEnum): + """Beam-search stopping mode, mirroring HuggingFace's tri-state + ``early_stopping`` (``True`` / ``False`` / ``"never"``). + + An ``IntEnum`` so it stays interchangeable with the raw integers coming from + ``sampling_config.early_stopping`` and used in the strategy grouping key. + Member names follow HF's ``early_stopping`` values (``True`` / ``False`` / + ``"never"``) directly. + + The stopping decision differs only in the upper bound used for the best + score an unfinished beam could still attain (score is + ``cum_log_prob / gen_length ** length_penalty``, and ``cum_log_prob <= 0``): + """ + + TRUE = 1 + """HF ``True`` (default): stop once ``beam_width`` finished candidates exist. + Regular (non-CBA) path.""" + + FALSE = 0 + """HF ``False``: CBA path bounding attainability by the beam's current + score.""" + + NEVER = 2 + """HF ``"never"``: CBA path bounding attainability by ``max_seq_len`` when + ``length_penalty > 0``. The CBA path treats any value other than ``0`` / + ``1`` as this mode.""" + + @classmethod + def from_raw(cls, value: Optional[int]) -> "BeamSearchEarlyStop": + """Map a raw ``sampling_config.early_stopping`` value to a mode. + + ``None`` -> ``TRUE`` (the default); ``0`` -> ``FALSE``; every other + integer -> ``NEVER`` (HF's "never"), matching the CBA path which + special-cases only ``FALSE``.""" + if value is None: + return cls.TRUE + if value == cls.FALSE: + return cls.FALSE + if value == cls.TRUE: + return cls.TRUE + return cls.NEVER + + +def _beam_topk(values: torch.Tensor, k: int) -> tuple[torch.Tensor, torch.Tensor]: + """Sorted top-k over the last dim of a 2D tensor, dispatched on row width. + + flashinfer's radix-select kernel is O(n) and much faster than torch.topk on + large rows (e.g. vocab-sized logits), while torch.topk wins on small rows + where the radix kernel's fixed per-call cost dominates. The 10k crossover + follows flashinfer's own guidance (``flashinfer.top_k`` docstring). + """ + if values.size(-1) > 10000: + return radix_topk_op(values, k) + return torch.topk(values, k=k, dim=-1, sorted=True) + + +@dataclass(kw_only=True) +class _CBAFields: + """The candidate-beams-array (CBA) tensors, shared by the persistent + ``BeamSearchStore`` and the per-step ``CBAState`` view (which reference the + same tensors). Field shapes and semantics are documented here once; both + subclasses inherit these fields. + + Only maintained for requests using exhaustive early_stopping modes + (early_stopping != TRUE); allocated lazily by + ``BeamSearchStore.ensure_cba`` on the first such request. + + Fields written on every beam-search step regardless of stopping mode + (``original_tokens``, ``prompt_lens``, ``batch_dones``) live on + ``BeamSearchStore`` instead, since they must exist before any CBA + request arrives. + """ + + cba_tokens: torch.Tensor + """[max_num_sequences, max_beam_width, max_seq_len] int32, finished-beam + path snapshots (generated tokens, BEAM_SEARCH_PAD_TOKEN padded).""" + cba_cum_log_probs: torch.Tensor + """[max_num_sequences, max_beam_width] float32, raw cumulative log-probs.""" + cba_normed_scores: torch.Tensor + """[max_num_sequences, max_beam_width] float32, length-normalized scores + (-inf: empty entry).""" + cba_lengths: torch.Tensor + """[max_num_sequences, max_beam_width] int32, generated lengths.""" + cba_caps: torch.Tensor + """[max_num_sequences] int32, per-slot CBA capacity (the request's + maximum beam width). Differs from the step's beam_width_out for + variable-beam-width requests.""" + original_log_probs: torch.Tensor + """[max_num_sequences, max_beam_width, max_seq_len] float32, uncorrected + per-slot sampled log-prob per step (log-prob analog of original_tokens). + Written and read by the CBA path for logprobs.""" + cba_log_probs: torch.Tensor + """[max_num_sequences, max_beam_width, max_seq_len] float32, per-token + log-probs of the CBA path snapshots.""" + + +@dataclass(kw_only=True) +class CBAState(_CBAFields): + """Candidate-Beams-Array (CBA) state, present only for requests using + exhaustive early_stopping modes (early_stopping != TRUE); see + beam_search_sampling_batch_cba. + + A per-step view over the persistent ``BeamSearchStore`` (shared CBA tensors + inherited from ``_CBAFields``) plus the step-local fields below. Bundling + them lets callers gate on a single ``cba is not None`` check (the enclosing + ``BeamSearchMetadata.cba`` is None outside CBA mode) instead of testing each + field. + """ + + end_ids: torch.Tensor + """[max_num_sequences] int32, per-slot end token id (< 0: no end token).""" + prompt_lens: torch.Tensor + """[max_num_sequences] int32, per-slot prompt length (from BeamSearchStore).""" + original_tokens: torch.Tensor + """[max_num_sequences, max_beam_width, max_seq_len] int32, uncorrected + per-slot tokens (from BeamSearchStore).""" + batch_dones: torch.Tensor + """[max_num_sequences] bool, per-slot termination verdict (from + BeamSearchStore).""" + max_seq_len: int = 0 + """Maximum sequence length (prompt + generated), used by the + best-attainable-score bound of the "never" early-stopping modes.""" + max_gen_len: int = 0 + """Host-known upper bound on the generated length (including this + step's token) across the group's requests. Bounds the width of the CBA + path-snapshot and pool-merge operations, which would otherwise run at + the full max_seq_len width every step. 0 means unknown (full width).""" + + +@dataclass(kw_only=True) +class BeamSearchStore: + """Persistent per-sampler beam-search storage. + + The candidate-beams-array tensors live in the optional ``cba`` member, + allocated by :meth:`ensure_cba` on the first request using an exhaustive + early_stopping mode. Beam search with the default ``early_stopping=TRUE`` + never touches them, so they are not allocated for it. + """ + + cache_indirection: torch.Tensor + """[max_num_sequences, max_beam_width, attention_size] int32, cache + indirection for beam search sampling.""" + cache_indirection_buffer: torch.Tensor + """[max_num_sequences, max_beam_width, attention_size] int32, second buffer + used to update the cache indirection during sampling.""" + cum_log_probs: torch.Tensor + """[max_num_sequences, max_beam_width] float32, current cumulative logprob + of each active beam.""" + first_finish_reasons: torch.Tensor + """[max_num_sequences, max_beam_width] int32, first finish reason per beam.""" + predecessor_beams: torch.Tensor + """[max_num_sequences, max_beam_width] int32, predecessor beam per beam, used + for stop word detection.""" + seq_offsets: torch.Tensor + """[max_num_sequences] int64, cached ``arange(max_num_sequences) * + max_beam_width`` used by ``beam_search_sampling_batch`` to flatten + (batch_idx, beam_idx) pairs.""" + beam_idx_arange: torch.Tensor + """[max_beam_width] int32, cached ``arange(max_beam_width)`` used as the + scatter source in the per-step ``cache_indirection.scatter_``.""" + beam_gen_lengths: torch.Tensor + """[max_num_sequences, max_beam_width] int32, number of generated tokens per + beam (frozen once a beam finishes). Only maintained for requests with a + non-zero beam-search length_penalty.""" + original_tokens: torch.Tensor + """[max_num_sequences, max_beam_width, max_seq_len] int32, uncorrected + per-slot tokens, written every beam-search step; read with + ``cache_indirection`` to snapshot finished paths into the CBA, and seeded + by the disaggregated first-generation handoff.""" + prompt_lens: torch.Tensor + """[max_num_sequences] int32, per-slot prompt length; used to derive + generated lengths and snapshot positions.""" + batch_dones: torch.Tensor + """[max_num_sequences] bool, per-slot beam-search termination verdict.""" + cba: Optional[_CBAFields] = None + """CBA tensors; None until the first exhaustive-early_stopping request.""" + + @classmethod + def create( + cls, + *, + cache_indirection_shape: tuple[int, ...], + max_num_sequences: int, + max_beam_width: int, + ) -> "BeamSearchStore": + """Allocate the per-sampler beam-search buffers. + + ``cache_indirection_shape`` is [max_num_sequences, max_beam_width, + attention_size]; the per-beam scalar buffers use its leading two dims. + """ + per_beam = cache_indirection_shape[:-1] + return cls( + cache_indirection=torch.empty(cache_indirection_shape, device="cuda", dtype=torch.int), + cache_indirection_buffer=int_tensor(cache_indirection_shape), + cum_log_probs=torch.empty(per_beam, device="cuda", dtype=torch.float32), + predecessor_beams=int_tensor(per_beam), + original_tokens=int_tensor(cache_indirection_shape), + first_finish_reasons=int_tensor(per_beam), + seq_offsets=( + torch.arange(max_num_sequences, device="cuda", dtype=torch.int64) * max_beam_width + ), + beam_idx_arange=torch.arange(max_beam_width, device="cuda", dtype=torch.int32), + beam_gen_lengths=int_tensor(per_beam), + prompt_lens=int_tensor((max_num_sequences,)), + batch_dones=torch.zeros((max_num_sequences,), device="cuda", dtype=torch.bool), + ) + + def ensure_cba(self) -> _CBAFields: + """Allocate the CBA tensors on first use and return them. + + Called when a request with an exhaustive early_stopping mode is + admitted; idempotent afterwards. + """ + if self.cba is None: + shape = self.original_tokens.shape + per_beam = shape[:-1] + num_sequences = shape[0] + self.cba = _CBAFields( + cba_tokens=int_tensor(tuple(shape)), + cba_cum_log_probs=torch.zeros(per_beam, device="cuda", dtype=torch.float32), + cba_normed_scores=torch.full( + per_beam, float("-inf"), device="cuda", dtype=torch.float32 + ), + cba_lengths=int_tensor(per_beam), + cba_caps=int_tensor((num_sequences,)), + original_log_probs=torch.zeros(tuple(shape), device="cuda", dtype=torch.float32), + cba_log_probs=torch.zeros(tuple(shape), device="cuda", dtype=torch.float32), + ) + return self.cba + + +@dataclass(kw_only=True) +class BeamHistory: + """Per-beam corrected tokens and log-probs. The three log-prob fields are + None unless log-probs are requested.""" + + tokens: torch.Tensor + """[num_beams, seq_len] int, corrected token ids per beam.""" + logprobs: torch.Tensor | None = None + """[num_beams, seq_len] float, per-token sampled log-prob.""" + logprobs_indices: torch.Tensor | None = None + """[num_beams, seq_len] int, vocab indices of the sampled log-probs.""" + cum_logprobs: torch.Tensor | None = None + """[num_beams] float, cumulative log-prob per beam.""" + + +@dataclass(kw_only=True, frozen=True) +class _BeamHistoryLogProbsSlices: + """Correlated beam-history log-prob tensors; all three fields are bound together.""" + + sampled_log_probs: torch.Tensor + sampled_logprobs_indices: torch.Tensor + cum_logprobs: torch.Tensor + + +@dataclass(kw_only=True, frozen=True) +class _BeamHistoryTensors: + """Beam-history tensor slices. + + Used to carry both device-side views (before D2H) and host-side + snapshots (after D2H). `log_probs` is bound iff log-probs are + requested. + """ + + cache_indirection: torch.Tensor + current_path: torch.Tensor + log_probs: _BeamHistoryLogProbsSlices | None + + +def _gather_beam_path( + *, current_path: torch.Tensor, cache_indirection: torch.Tensor +) -> torch.Tensor: + """Gather the correct tokens for each beam from current_path.""" + new_path = torch.zeros_like(current_path) + torch.gather(input=current_path, dim=0, index=cache_indirection, out=new_path) + return new_path + + +@dataclass(kw_only=True) +class BeamSearchMetadata(StrategyMetadata): + """Stateful tensors required by beam_search_sampling_batch.""" + + cache_indirection: torch.Tensor + cache_indirection_buffer: torch.Tensor + cum_log_probs: torch.Tensor + new_log_probs: torch.Tensor + seq_slots: torch.Tensor + seq_lens: torch.Tensor + finished_beams: torch.Tensor + predecessor_beams: torch.Tensor + seq_offsets: torch.Tensor + beam_idx_arange: torch.Tensor + beam_gen_lengths: torch.Tensor + stop_past_tokens: Optional[torch.Tensor] = None + """[max_stop_word_length, max_num_sequences, max_beam_width] int32, the + finish handler's rolling stop-word window (FinishReasonsHandler store). + Used by both the regular and CBA paths: the beam axis is reordered by the + step's predecessor beams so the handler's stop-word matching stays correct + across beam swaps. When None (tests without stop words), the reorder is + skipped — multi-token stop-word matching would then be unreliable across + beam swaps.""" + cba: Optional[CBAState] = None + """Candidate-Beams-Array state, present only for exhaustive early_stopping + modes (early_stopping != TRUE); None for the regular beam-search path.""" + + +def _update_cache_indirection_buffer( + cache_indirection_input: torch.Tensor, + cache_indirection_output: torch.Tensor, + seq_slots: torch.Tensor, +) -> None: + assert cache_indirection_input.device == cache_indirection_output.device + cache_indirection_input.index_copy_(0, seq_slots, cache_indirection_output[seq_slots]) + + +def _beam_step_preprocess( + logits: torch.Tensor, + *, + beam_width_in: int, + temperature: float | None, + return_probs: bool, + args: "BeamSearchMetadata", +) -> tuple[torch.Tensor, Optional[torch.Tensor], int]: + """Shared front-end of both beam-search step ops. + + Applies temperature, snapshots the cache indirection into its buffer, and + returns ``(logprobs, softmax, batch_size)``. ``softmax`` is None when + ``return_probs`` is False. + """ + assert logits.dim() == 2, "logits should be 2D: [batch_size * beam_width, vocab_size]" + batch_size, vocab_size = logits.size() + batch_size = batch_size // beam_width_in + + logits = logits.view(batch_size, beam_width_in, vocab_size) + if temperature is not None and temperature != 0: + logits = logits / max(temperature, 1e-5) + softmax: Optional[torch.Tensor] = None + if return_probs: + softmax = torch.softmax(logits, dim=-1) + _update_cache_indirection_buffer( + args.cache_indirection_buffer, args.cache_indirection, args.seq_slots + ) + assert batch_size == args.seq_slots.size(0) + + return torch.log_softmax(logits, dim=-1), softmax, batch_size + + +def _pad_next_tokens(next_tokens: torch.Tensor, store_width: int) -> torch.Tensor: + """Pad a [batch, beam_width_out] token tensor to the store's beam width. + + The batched sampling buffers are allocated at the maximum beam width; on + variable-beam-width steps the op produces fewer columns, and the padding + (BEAM_SEARCH_PAD_TOKEN) is never consumed — finalization rewrites all + beam tokens from the corrected paths. + """ + if next_tokens.size(1) >= store_width: + return next_tokens + return torch.nn.functional.pad( + next_tokens, (0, store_width - next_tokens.size(1)), value=BEAM_SEARCH_PAD_TOKEN + ) + + +def beam_candidate_topk( + logprobs: torch.Tensor, + *, + beam_width_out: int, + length_penalty: "torch.Tensor | float | None" = None, + cand_gen_lengths: Optional[torch.Tensor] = None, + diversity_rate: "torch.Tensor | float | None" = None, + source_beam_indices: Optional[torch.Tensor] = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Two-stage top-k over beam-expansion candidates with per-source-beam + ranking adjustments, ranked by:: + + (cum_log_prob + diversity_rate * source_beam_index) + / gen_length**length_penalty + + The diversity term ``diversity_rate * source_beam_index`` spreads selection + across source beams (candidates from lower-ranked beams get a boost), and + the length term normalizes by ``gen_length**length_penalty``. Finished + beams stay frozen in their slots rather than in a separate pool, so a + finished beam's (single) candidate receives the same ``rate * slot_index`` + boost as any other — slot position thus slightly affects how hard a + finished beam is to evict. + + Mathematically equivalent to adjusting the full [batch, bw_in, vocab] + candidate matrix and taking a flat top-k, but avoids touching the vocab + axis: both adjustments are constant per source beam, so they cannot + change the ordering *within* a beam. Any global winner is therefore + among its own beam's top-``beam_width_out`` raw candidates. Stage 1 + takes a per-beam top-k on raw scores; stage 2 adjusts only the + ``bw_in * bw_out`` survivors and selects the global top-k. + + Args: + logprobs: [batch, beam_width_in, vocab] raw cumulative log-probs. + length_penalty: scalar, or per-request tensor of shape [batch]. + None/0 disables length normalization. + cand_gen_lengths: [batch, beam_width_in] candidate generated lengths. + Required iff ``length_penalty`` is active. + diversity_rate: scalar, or per-request tensor of shape [batch]. + None/0 disables the diversity adjustment. + source_beam_indices: optional cached ``arange(beam_width_in)`` on the + logprobs device (e.g. ``BeamSearchStore.beam_idx_arange``), used + by the diversity adjustment; computed on the fly if omitted. + + Returns: + (sorted_logprobs, predecessor_beams, tokens): the raw (unadjusted) + cumulative log-probs of the selected candidates, the source-beam index + each candidate expands from, and its token id — all of shape + [batch, beam_width_out] (indices int32), ordered by descending + adjusted score. + """ + batch_size, beam_width_in, vocab_size = logprobs.shape + # Clamps are only relevant for tiny test vocabularies: stage 1 cannot + # exceed the vocab, the global top-k cannot exceed the pooled candidates. + stage1_k = min(beam_width_out, vocab_size) + beam_width_out = min(beam_width_out, beam_width_in * stage1_k) + # Stage 1: raw per-beam top-k (the per-beam adjustments are constant + # along the vocab axis, so raw ordering == adjusted ordering). + per_beam_vals, per_beam_tokens = _beam_topk( + logprobs.view(batch_size * beam_width_in, vocab_size), stage1_k + ) + per_beam_vals = per_beam_vals.view(batch_size, beam_width_in, stage1_k) + per_beam_tokens = per_beam_tokens.view(batch_size, beam_width_in, stage1_k) + # Stage 2: adjust only the survivors and pick the global top-k. + keys = per_beam_vals + if diversity_rate is not None: + rate = ( + diversity_rate.view(-1, 1, 1) + if isinstance(diversity_rate, torch.Tensor) + else diversity_rate + ) + if source_beam_indices is None: + source_beam_indices = torch.arange( + beam_width_in, device=logprobs.device, dtype=torch.int32 + ) + keys = keys + rate * source_beam_indices[:beam_width_in].view(1, -1, 1) + if length_penalty is not None: + assert cand_gen_lengths is not None, ( + "cand_gen_lengths is required when length_penalty is active" + ) + exponent = ( + length_penalty.view(-1, 1) + if isinstance(length_penalty, torch.Tensor) + else length_penalty + ) + # Candidate lengths are >= 1 by construction (active beams: + # generated + 1; finished beams froze after generating at least one + # token), so the power is always well-defined. + penalty_factor = cand_gen_lengths.to(logprobs.dtype).pow(exponent) + keys = keys / penalty_factor.unsqueeze(-1) + _, selected = _beam_topk(keys.view(batch_size, -1), beam_width_out) + sorted_logprobs = per_beam_vals.view(batch_size, -1).gather(1, selected) + predecessor_beams = (selected // stage1_k).to(torch.int32) + tokens = per_beam_tokens.view(batch_size, -1).gather(1, selected).to(torch.int32) + return sorted_logprobs, predecessor_beams, tokens + + +def beam_search_sampling_batch( + logits: torch.Tensor, + *, + beam_width_in: int, + beam_width_out: int, + beam_search_args: BeamSearchMetadata, + temperature: float | None, + length_penalty: "torch.Tensor | float | None" = None, + diversity_rate: "torch.Tensor | float | None" = None, + return_probs: bool = True, +) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """Sample beam_width tokens for each request in parallel. + + ``length_penalty`` normalizes the beam-selection ranking key as + ``cum_log_prob / gen_length**length_penalty``, and ``diversity_rate`` adds + ``diversity_rate * source_beam_index`` to it; see ``beam_candidate_topk``. The stored ``cum_log_probs`` remain + raw. Both accept a per-request tensor of shape [batch_size] or a scalar; + None/0 disables the respective adjustment. When ``length_penalty`` is + active, per-beam generated lengths are maintained in place in + ``beam_search_args.beam_gen_lengths`` (alongside the other stateful + metadata tensors this function updates). + """ + logprobs, softmax, batch_size = _beam_step_preprocess( + logits, + beam_width_in=beam_width_in, + temperature=temperature, + return_probs=return_probs, + args=beam_search_args, + ) + + finished_beams_mask = ( + beam_search_args.finished_beams[beam_search_args.seq_slots, :beam_width_in] + != FinishReason.NOT_FINISHED.value + ) + finished_beams_mask_expanded = finished_beams_mask.unsqueeze(-1).expand( + -1, -1, logprobs.size(-1) + ) + logprobs = torch.where(finished_beams_mask_expanded, float("-inf"), logprobs) + logprobs[..., 0] = torch.where(finished_beams_mask, 0, logprobs[..., 0]) + + logprobs += beam_search_args.cum_log_probs.unsqueeze(-1)[ + beam_search_args.seq_slots, :beam_width_in + ] + + if not isinstance(length_penalty, torch.Tensor) and not length_penalty: + length_penalty = None # scalar 0 (or None) disables normalization + if not isinstance(diversity_rate, torch.Tensor) and not diversity_rate: + diversity_rate = None # scalar 0 (or None) disables the adjustment + cand_gen_lengths: Optional[torch.Tensor] = None + if length_penalty is not None: + # Candidate generated length: active beams grow by one token this + # step, finished beams keep their frozen length (they only append + # pads). + # Candidate generated length: active beams grow by one token this + # step, finished beams keep their frozen length (they only append + # pads). The counter is per-beam and cannot be derived from + # seq_len - prompt_len, which is shared by all beams of a request. + gen_lengths = beam_search_args.beam_gen_lengths[beam_search_args.seq_slots, :beam_width_in] + cand_gen_lengths = gen_lengths + (~finished_beams_mask).to(gen_lengths.dtype) + # Rank by the (optionally adjusted) score; keep raw cum_log_probs for + # storage. The two-stage selection is used even without adjustments: it is + # equivalent to a flat top-k there, and faster with the radix backend + # (more, shorter rows parallelize better). + sorted_logprobs, predecessor_beam, next_tokens = beam_candidate_topk( + logprobs, + beam_width_out=beam_width_out, + length_penalty=length_penalty, + cand_gen_lengths=cand_gen_lengths, + diversity_rate=diversity_rate, + source_beam_indices=beam_search_args.beam_idx_arange, + ) + + if cand_gen_lengths is not None: + # Zero the full configured width first: on a narrowing + # variable-beam-width step beam_width_out < beam_width_in, and the + # stop criterion reads the full width, so the untouched tail would + # otherwise keep the previous step's lengths. + beam_search_args.beam_gen_lengths[beam_search_args.seq_slots] = 0 + beam_search_args.beam_gen_lengths[beam_search_args.seq_slots, :beam_width_out] = ( + torch.gather(cand_gen_lengths, dim=1, index=predecessor_beam) + ) + beam_search_args.predecessor_beams[beam_search_args.seq_slots, :beam_width_out] = ( + predecessor_beam + ) + if beam_search_args.stop_past_tokens is not None: + # Reorder the finish handler's rolling stop-word window to follow the + # beam swap, so multi-token stop-word matching (which appends this + # step's tokens to the window after this op) compares against the + # correct per-beam history. + window = beam_search_args.stop_past_tokens[:, beam_search_args.seq_slots, :beam_width_in] + beam_search_args.stop_past_tokens[:, beam_search_args.seq_slots, :beam_width_out] = ( + torch.gather( + window, + 2, + predecessor_beam.long().unsqueeze(0).expand(window.size(0), -1, -1), + ) + ) + + finished_beams = beam_search_args.finished_beams[beam_search_args.seq_slots].view(-1) + + offset_predecessor_beam = predecessor_beam + beam_search_args.seq_offsets[ + : predecessor_beam.size(0) + ].unsqueeze(1) + finished_beams = finished_beams[offset_predecessor_beam] + # Write only the beam_width_out columns produced this step: with variable + # beam width it can be smaller than the store width (the stale view() of + # the full width crashed on such steps). + beam_search_args.finished_beams[beam_search_args.seq_slots, :beam_width_out] = finished_beams + + cache_indirection = beam_search_args.cache_indirection[ + beam_search_args.seq_slots, :beam_width_out + ] + cache_indirection_buffer = beam_search_args.cache_indirection_buffer[ + beam_search_args.seq_slots, :beam_width_in + ] + torch.gather( + cache_indirection_buffer, + dim=1, + index=predecessor_beam.unsqueeze(2).expand(-1, -1, cache_indirection.size(2)), + out=cache_indirection, + ) + + index = beam_search_args.seq_lens.view(-1, 1, 1).expand(-1, beam_width_out, 1) + src = ( + beam_search_args.beam_idx_arange[:beam_width_out] + .view(1, beam_width_out, 1) + .expand(batch_size, beam_width_out, 1) + ) + cache_indirection.scatter_(2, index, src) + + beam_search_args.cache_indirection[beam_search_args.seq_slots, :beam_width_out] = ( + cache_indirection + ) + + ended_predecessor_mask = torch.gather(dim=1, index=predecessor_beam, input=finished_beams_mask) + next_tokens = torch.where(ended_predecessor_mask, BEAM_SEARCH_PAD_TOKEN, next_tokens) + + old_cum_log_probs = beam_search_args.cum_log_probs[beam_search_args.seq_slots].view(-1) + beam_search_args.new_log_probs[beam_search_args.seq_slots, :beam_width_out] = ( + sorted_logprobs[:, :beam_width_out] - old_cum_log_probs[offset_predecessor_beam] + ) + beam_search_args.cum_log_probs[beam_search_args.seq_slots, :beam_width_out] = sorted_logprobs[ + :, :beam_width_out + ] + return _pad_next_tokens(next_tokens, beam_search_args.finished_beams.size(1)), softmax + + +class CBAStepResult(NamedTuple): + """Return of ``_cba_step_math``. A NamedTuple (not a dataclass) so it is a + valid output of the torch.compile'd fullgraph function while still naming + the eleven tensors the caller writes back.""" + + slot_pred: torch.Tensor + """[bs, num_beams] int32, predecessor beam of each continuing slot.""" + slot_tok: torch.Tensor + """[bs, num_beams] int32, token of each continuing slot.""" + slot_cum: torch.Tensor + """[bs, num_beams] float32, cumulative log-prob of each continuing slot.""" + step_log_probs: torch.Tensor + """[bs, num_beams] float32, this step's per-slot log-prob.""" + top_normed: torch.Tensor + """[bs, pool_width] float32, merged CBA pool normalized scores.""" + merged_cum: torch.Tensor + """[bs, pool_width] float32, merged CBA pool cumulative log-probs.""" + merged_len: torch.Tensor + """[bs, pool_width] int32, merged CBA pool lengths.""" + merged_tokens: torch.Tensor + """[bs, pool_width, snap_len] int32, merged CBA pool path snapshots.""" + merged_lps: torch.Tensor + """[bs, pool_width, snap_len] float32, merged CBA pool per-token log-probs.""" + done: torch.Tensor + """[bs] bool, per-slot beam-search termination verdict.""" + reordered_window: Optional[torch.Tensor] + """Stop-word window reordered to follow the beam swap, or None when no + stop-word window is present.""" + + +def _cba_step_math( + # candidates (from beam_candidate_topk) + cand_cum: torch.Tensor, # [bs, C] float32 + cand_pred: torch.Tensor, # [bs, C] int32 + cand_tok: torch.Tensor, # [bs, C] int32 + # per-request state (store tensors + the group's slots) + slots: torch.Tensor, # [bs] int64 + seq_lens: torch.Tensor, # [bs] int32 + snap_arange: torch.Tensor, # [S] int64, S = bounded snapshot width + exponent: torch.Tensor, # [bs, 1] float32, 0 == no length penalty + cache_indirection: torch.Tensor, + original_tokens: torch.Tensor, + original_log_probs: torch.Tensor, + cum_log_probs: torch.Tensor, + finished_beams: torch.Tensor, + end_ids: torch.Tensor, + prompt_lens: torch.Tensor, + cba_caps: torch.Tensor, + cba_normed_scores: torch.Tensor, + cba_cum_log_probs: torch.Tensor, + cba_lengths: torch.Tensor, + cba_tokens: torch.Tensor, + cba_log_probs: torch.Tensor, + stop_past_tokens: Optional[torch.Tensor], + # static (graph-specializing) parameters + beam_width_in: int, + num_beams: int, + early_stopping: int, # BeamSearchEarlyStop; only ``FALSE`` (0) is special-cased + max_seq_len: int, +) -> CBAStepResult: + """Small-tensor math of the CBA beam-search step (see + beam_search_sampling_batch_cba), written without data-dependent shapes so + it can be fused by torch.compile: everything between candidate selection + and the store writebacks. Pure — all mutations happen in the caller. + + NB: this function is torch.compile'd (see ``_cba_step_compiled``). Keep it + free of data-dependent shapes and of in-place ops on the input tensors + (use out-of-place ``masked_fill`` etc.), so tracing stays fullgraph-clean. + """ + batch_size, num_candidates = cand_cum.shape + neg_inf = float("-inf") + + harvest_mask = finished_beams[slots, :beam_width_in] != FinishReason.NOT_FINISHED.value + end_ids_b = end_ids[slots].view(-1, 1) + caps = cba_caps[slots].view(-1, 1) + prompts = prompt_lens[slots].view(-1, 1) + gen_lens = (seq_lens - prompts.view(-1)).view(-1, 1) + cand_len = gen_lens + 1 + + # -inf candidates (from harvested rows) must count as slot fillers, not + # end candidates, to preserve the >= K actives invariant even when every + # beam was harvested at once. + is_end = (cand_tok == end_ids_b) & torch.isfinite(cand_cum) + cand_rank = torch.arange(num_candidates, device=cand_cum.device).view(1, -1) + + # --- Beam slots continue with the first K non-end candidates. Scatter + # through a K+1-wide buffer instead of masked_select (data-dependent + # shapes cannot be compiled); non-selected entries all collide on the + # spare column and are discarded. + active_mask = ~is_end + active_pos = torch.cumsum(active_mask.to(torch.int32), dim=1) - 1 + scatter_idx = torch.where(active_mask & (active_pos < num_beams), active_pos, num_beams).long() + + def _take(src: torch.Tensor) -> torch.Tensor: + buf = src.new_zeros(batch_size, num_beams + 1) + buf.scatter_(1, scatter_idx, src) + return buf[:, :num_beams] + + slot_pred = _take(cand_pred) + slot_tok = _take(cand_tok) + slot_cum = _take(cand_cum) + step_log_probs = slot_cum - cum_log_probs[slots, :beam_width_in].gather(1, slot_pred.long()) + + # --- CBA insertion: normalized scores of the new end-token candidates. + new_normed = cand_cum / cand_len.to(cand_cum.dtype).pow(exponent) + eligible = is_end & (cand_rank < caps) + new_normed = new_normed.masked_fill(~eligible, neg_inf) + + # Snapshot candidate paths through the cache indirection (the work tree + # is rewritten by later steps). Indirection entries beyond the current + # length are uninitialized: lanes are masked below, but gather indices + # must be clamped in-bounds first. + step_idx = (prompts + snap_arange.view(1, -1)).clamp(max=cache_indirection.size(-1) - 1) + step_idx_e = step_idx.unsqueeze(1).expand(-1, beam_width_in, -1) + ind_at = torch.gather(cache_indirection[slots, :beam_width_in].long(), 2, step_idx_e) + ind_at = ind_at.clamp_(0, beam_width_in - 1) + tok_at = torch.gather(original_tokens[slots, :beam_width_in], 2, step_idx_e) + lp_at = torch.gather(original_log_probs[slots, :beam_width_in], 2, step_idx_e) + snap_len = snap_arange.size(0) + parent_exp = cand_pred.long().unsqueeze(-1).expand(-1, -1, snap_len) + src_beam = torch.gather(ind_at, 1, parent_exp) + t_valid = snap_arange.view(1, 1, -1) < gen_lens.view(-1, 1, 1) + new_paths = torch.gather(tok_at, 1, src_beam).masked_fill(~t_valid, BEAM_SEARCH_PAD_TOKEN) + new_lp_paths = torch.gather(lp_at, 1, src_beam).masked_fill(~t_valid, 0.0) + end_pos = gen_lens.view(-1, 1, 1).expand(-1, num_candidates, 1).clamp(max=snap_len - 1) + new_paths = new_paths.scatter(2, end_pos, cand_tok.unsqueeze(-1).to(new_paths.dtype)) + # the terminating token's own log-prob = candidate cum - parent cum + parent_cum = cum_log_probs[slots, :beam_width_in].gather(1, cand_pred.long()) + new_lp_paths = new_lp_paths.scatter(2, end_pos, (cand_cum - parent_cum).unsqueeze(-1)) + + # Harvested beams (stop-word finishes latched after the previous step): + # their recorded tokens already include the terminating stop word, so the + # snapshot is the beam's own path (parent = itself) at the current length. + harvest_paths = torch.gather(tok_at, 1, ind_at).masked_fill(~t_valid, BEAM_SEARCH_PAD_TOKEN) + harvest_lp_paths = torch.gather(lp_at, 1, ind_at).masked_fill(~t_valid, 0.0) + harvest_cum = cum_log_probs[slots, :beam_width_in] + harvest_normed = harvest_cum / gen_lens.to(harvest_cum.dtype).pow(exponent) + harvest_normed = harvest_normed.masked_fill(~harvest_mask, neg_inf) + + # --- Pool merge: keep the best pool_width by normalized score, then + # enforce the per-request capacity (replace-min against the pool). + pool_width = cba_normed_scores.size(1) + all_normed = torch.cat([cba_normed_scores[slots], new_normed, harvest_normed], dim=1) + top_normed, top_i = torch.topk(all_normed, k=pool_width, sorted=True, dim=-1) + top_normed = top_normed.masked_fill( + torch.arange(pool_width, device=cand_cum.device).view(1, -1) >= caps, neg_inf + ) + all_cum = torch.cat([cba_cum_log_probs[slots], cand_cum, harvest_cum], dim=1) + all_len = torch.cat( + [ + cba_lengths[slots], + cand_len.expand(-1, num_candidates), + gen_lens.expand(-1, beam_width_in), + ], + dim=1, + ) + all_tokens = torch.cat([cba_tokens[slots, :, :snap_len], new_paths, harvest_paths], dim=1) + all_lps = torch.cat([cba_log_probs[slots, :, :snap_len], new_lp_paths, harvest_lp_paths], dim=1) + merged_cum = all_cum.gather(1, top_i) + merged_len = all_len.gather(1, top_i) + top_i_wide = top_i.unsqueeze(-1).expand(-1, -1, snap_len) + merged_tokens = all_tokens.gather(1, top_i_wide) + merged_lps = all_lps.gather(1, top_i_wide) + + # --- Done verdict: CBA full, and the best candidate's + # attainable normalized score cannot beat the worst kept entry. + # early_stopping != FALSE (i.e. HF "never") assumes scores can still + # increase with length for positive penalties, so unfinished beams have no + # upper bound on attainability (they are bounded by max length); otherwise + # (FALSE) the beam's current score is the correct bound (assume scores + # decrease monotonically with sequence length, so longer sequences only get + # less attractive). + min_kept = top_normed.gather(1, (caps - 1).long()).view(-1) + if early_stopping != BeamSearchEarlyStop.FALSE: + max_gen = (max_seq_len - prompts.view(-1)).to(cand_len.dtype) + bound_len = torch.where(exponent.view(-1) > 0, max_gen, cand_len.view(-1)) + else: + bound_len = cand_len.view(-1) + best_attainable = cand_cum[:, 0] / bound_len.to(cand_cum.dtype).pow(exponent.view(-1)) + done = (min_kept > neg_inf) & (min_kept >= best_attainable) + + # Reorder the finish handler's rolling stop-word window to follow the + # beam swap (matching stays correct across swaps). + reordered_window = None + if stop_past_tokens is not None: + window = stop_past_tokens[:, slots, :beam_width_in] + reordered_window = torch.gather( + window, 2, slot_pred.long().unsqueeze(0).expand(window.size(0), -1, -1) + ) + + return CBAStepResult( + slot_pred=slot_pred, + slot_tok=slot_tok, + slot_cum=slot_cum, + step_log_probs=step_log_probs, + top_normed=top_normed, + merged_cum=merged_cum, + merged_len=merged_len, + merged_tokens=merged_tokens, + merged_lps=merged_lps, + done=done, + reordered_window=reordered_window, + ) + + +# Compiled lazily on first use; the first CBA-mode request of a process pays +# the inductor compile once (subsequent shapes are covered by the dynamic +# batch/snapshot-length dims marked at the call site). +_cba_step_compiled = torch.compile(_cba_step_math, dynamic=None, fullgraph=True) + + +def beam_search_sampling_batch_cba( + logits: torch.Tensor, + *, + beam_width_in: int, + beam_width_out: int, + beam_search_args: BeamSearchMetadata, + temperature: float | None, + early_stopping: int, # BeamSearchEarlyStop + length_penalty: "torch.Tensor | float | None" = None, + diversity_rate: "torch.Tensor | float | None" = None, + return_probs: bool = True, +) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """Beam-search step with a candidate-beams array (CBA) for the exhaustive + early_stopping modes (``early_stopping != TRUE``): + + - The top ``2 * beam_width`` expansion candidates (ranked by raw + cumulative log-prob, plus the optional diversity adjustment — the + length penalty does NOT enter candidate ranking) are split by end-token: + end-token candidates ranked within the top ``beam_width`` are inserted + into the CBA, which keeps the best ``beam_width`` finished paths seen so + far by length-normalized score (path snapshots are taken eagerly since + the work tree is rewritten by later steps). All beam slots then continue + with the best non-end candidates, so exploration never narrows. + - Stop words (any length) are detected by the finish handler after the + step, which latches a per-beam finish reason; at the START of the next + step this op harvests the latched beams into the CBA (their paths are + complete, including the stop word) and masks their rows so the freed + slots refill with active candidates — equivalent to coercing a finished + beam into a top-ranked end-token candidate. + - A request is done when the CBA is full and the best active candidate's + attainable normalized score cannot beat the worst CBA entry. + ``early_stopping == FALSE`` bounds attainability by the beam's current score + (assume scores decrease monotonically with sequence length); any other + value places no upper bound on attainability for unfinished beams (assume + scores can increase with length, e.g. when ``length_penalty > 0``), so it + bounds by ``max_seq_len`` (HF's "never"). The verdict is published by + marking every beam slot finished, which drives the regular stop + machinery. + + Requires the CBA fields of ``BeamSearchMetadata`` to be set. + """ + args = beam_search_args + cba = args.cba + assert cba is not None, "CBA metadata is required for early_stopping != TRUE" + num_beams = beam_width_out + device = logits.device + slots = args.seq_slots + + logprobs, softmax, batch_size = _beam_step_preprocess( + logits, + beam_width_in=beam_width_in, + temperature=temperature, + return_probs=return_probs, + args=args, + ) + # Beams latched finished by the finish handler after the previous step + # are harvested into the CBA below; mask their rows so the freed slots + # refill with active candidates from the other beams. Only STOP_WORDS + # latches can reach this point: LENGTH fires on all beams at once (their + # lengths are uniform here) and the END_ID flood below is all-beams too — + # either way the request stops that same step, so there is no next step + # to harvest them (finalize handles those paths instead). + harvest_mask = args.finished_beams[slots, :beam_width_in] != FinishReason.NOT_FINISHED.value + logprobs = logprobs.masked_fill(harvest_mask.unsqueeze(-1), float("-inf")) + logprobs += args.cum_log_probs.unsqueeze(-1)[slots, :beam_width_in] + + # --- Top 2K candidates by raw score (+ diversity), two-stage. The + # "at least K non-end candidates" invariant that keeps the K beam slots + # fillable relies on each source beam contributing at most one *end-token* + # candidate, which holds because end-token detection matches the single + # per-request end id (``end_ids``) below. Other ways a beam can finish do + # NOT add end-token candidates and so cannot break the invariant: + # * Stop words (any length, including a second EOS supplied as a stop + # word) are latched by the finish handler at the previous step and + # harvested here (see the docstring). Their rows are masked to -inf + # above, so the isfinite guard on ``is_end`` counts them as active + # fillers, not end candidates. + # * LENGTH / all-beam END_ID floods stop the whole request that step, so + # there is no next step whose slots would need filling. + # A model with multiple EOS tokens is therefore supported by registering + # the extra EOS ids as stop words; matching several ids directly in + # ``is_end`` would instead require widening the candidate pool and is out + # of scope here. + if not isinstance(diversity_rate, torch.Tensor) and not diversity_rate: + diversity_rate = None + # Each source beam contributes at most one finite end-token candidate, so + # beam_width_in extra candidates on top of the slots to fill always leave + # >= num_beams active ones — also on narrowing variable-beam-width steps + # (beam_width_in > num_beams), where 2 * num_beams would not. + cand_cum, cand_pred, cand_tok = beam_candidate_topk( + logprobs, + beam_width_out=num_beams + max(num_beams, beam_width_in), + diversity_rate=diversity_rate, + source_beam_indices=args.beam_idx_arange, + ) + + # Length penalty normalized to a per-request exponent tensor (0 == off), + # so the compiled step math is branch-free. + if isinstance(length_penalty, torch.Tensor): + exponent = length_penalty.view(-1, 1).to(torch.float32) + else: + exponent = torch.full( + (batch_size, 1), float(length_penalty or 0.0), dtype=torch.float32, device=device + ) + + snap_len = cba.cba_tokens.size(-1) + if cba.max_gen_len > 0: + # Snapshots and pool merges only ever touch generated positions, and + # pool entry lengths are bounded by the running maximum generation + # length; columns beyond it keep their (unread) previous content. + snap_len = min(snap_len, cba.max_gen_len) + snap_arange = torch.arange(snap_len, device=device) + + # CPU callers (unit tests) take the eager function: inductor's CPU + # compile latency would dominate, and the fusion only pays off on CUDA. + step_fn = _cba_step_compiled if logits.is_cuda else _cba_step_math + if logits.is_cuda: + for t in (cand_cum, cand_pred, cand_tok, slots, args.seq_lens, exponent): + torch._dynamo.mark_dynamic(t, 0) + # snap_arange's length is the (small, bounded) snapshot width, which the + # compiled step reads as a size; newer dynamo raises if a mark_dynamic + # dim gets specialized to a constant, so allow specialization here. + torch._dynamo.maybe_mark_dynamic(snap_arange, 0) + ( + slot_pred, + slot_tok, + slot_cum, + step_log_probs, + top_normed, + merged_cum, + merged_len, + merged_tokens, + merged_lps, + done, + reordered_window, + ) = step_fn( + cand_cum=cand_cum, + cand_pred=cand_pred, + cand_tok=cand_tok, + slots=slots, + seq_lens=args.seq_lens, + snap_arange=snap_arange, + exponent=exponent, + cache_indirection=args.cache_indirection, + original_tokens=cba.original_tokens, + original_log_probs=cba.original_log_probs, + cum_log_probs=args.cum_log_probs, + finished_beams=args.finished_beams, + end_ids=cba.end_ids, + prompt_lens=cba.prompt_lens, + cba_caps=cba.cba_caps, + cba_normed_scores=cba.cba_normed_scores, + cba_cum_log_probs=cba.cba_cum_log_probs, + cba_lengths=cba.cba_lengths, + cba_tokens=cba.cba_tokens, + cba_log_probs=cba.cba_log_probs, + stop_past_tokens=args.stop_past_tokens, + beam_width_in=beam_width_in, + num_beams=num_beams, + early_stopping=early_stopping, + max_seq_len=cba.max_seq_len, + ) + # --- Writebacks (kept eager: mixed advanced/basic indexing on the store + # tensors, and the compiled math stays pure). + cba.cba_normed_scores[slots] = top_normed + cba.cba_cum_log_probs[slots] = merged_cum + cba.cba_lengths[slots] = merged_len + cba.cba_tokens[slots, :, :snap_len] = merged_tokens + cba.cba_log_probs[slots, :, :snap_len] = merged_lps + cba.batch_dones[slots] = done + # NOT_FINISHED == 0, so the verdict maps directly onto the finish reason. + # Flood the full row: the stop criterion reads the first py_beam_width + # (== capacity) entries, which can exceed this step's beam_width_out for + # variable-beam-width requests. + args.finished_beams[slots] = ( + done.view(-1, 1).to(torch.int32) * FinishReason.END_ID.value + ).expand(-1, args.finished_beams.size(1)) + stop_window = args.stop_past_tokens + if reordered_window is not None and stop_window is not None: + stop_window[:, slots, :num_beams] = reordered_window + + # --- Beam-slot state updates (same contract as beam_search_sampling_batch). + args.predecessor_beams[slots, :num_beams] = slot_pred + cache_indirection = args.cache_indirection[slots, :num_beams] + cache_indirection_buffer = args.cache_indirection_buffer[slots, :beam_width_in] + torch.gather( + cache_indirection_buffer, + dim=1, + index=slot_pred.long().unsqueeze(2).expand(-1, -1, cache_indirection.size(2)), + out=cache_indirection, + ) + index = args.seq_lens.view(-1, 1, 1).expand(-1, num_beams, 1) + src = args.beam_idx_arange[:num_beams].view(1, num_beams, 1).expand(batch_size, num_beams, 1) + cache_indirection.scatter_(2, index, src) + args.cache_indirection[slots, :num_beams] = cache_indirection + + args.new_log_probs[slots, :num_beams] = step_log_probs + # Record this step's per-slot log-prob at the emission position so path + # snapshots can recover per-token log-probs (analog of original_tokens). + olp = cba.original_log_probs[slots, :num_beams] + olp.scatter_( + 2, + args.seq_lens.view(-1, 1, 1).expand(-1, num_beams, 1).long(), + step_log_probs.unsqueeze(-1), + ) + cba.original_log_probs[slots, :num_beams] = olp + args.cum_log_probs[slots, :num_beams] = slot_cum + return _pad_next_tokens(slot_tok, args.finished_beams.size(1)), softmax + + +BeamHistoryBuilder: TypeAlias = Callable[[], BeamHistory | None] +"""Builder for BeamHistory. + +Used to defer possibly unnecessary host-tensor construction until update_requests(). +""" + + +@dataclass(kw_only=True) +class CBAGroupHost: + """Host-side snapshot of the CBA state for a group of beam-search requests, + produced by ``TorchSampler._prepare_cba_group_host`` and consumed by + ``_prepare_beam_history_cba``. One batched D2H copy per tensor covers the + whole group; the consumer slices per-request rows via ``pos[slot]``. + """ + + pos: dict[int, int] + """Maps a request's seq slot to its row index in the batched tensors below.""" + should_stop: torch.Tensor + cache_indirection: torch.Tensor + original_tokens: torch.Tensor + cum: torch.Tensor + cba_tokens: torch.Tensor + cba_cum: torch.Tensor + cba_normed: torch.Tensor + cba_lengths: torch.Tensor + original_log_probs: Optional[torch.Tensor] + """None when no request in the group requests log-probs.""" + cba_log_probs: Optional[torch.Tensor] + """None when no request in the group requests log-probs.""" + + +def _prepare_beam_search( + beam_search_store: BeamSearchStore, + log_probs_store: LogProbsStore, + seq_slots_long: torch.Tensor, + max_prompt_len: int, + prompt_lens_cuda: torch.Tensor, + beam_caps_cuda: torch.Tensor, +) -> None: + """Prepare the beam search buffers for the requests + + If the last context chunk is being processed, + initialize/reset the buffers for the request. + + ``seq_slots_long`` must be int64 (required by ``index_fill_``). + """ + beam_search_store.cache_indirection.narrow(2, 0, max_prompt_len).index_fill_( + 0, seq_slots_long, 0 + ) + beam_search_store.cum_log_probs.index_fill_(0, seq_slots_long, 0) + log_probs_store.sampled_log_probs.index_fill_(0, seq_slots_long, 0) + log_probs_store.sampled_log_prob_ranks.index_fill_(0, seq_slots_long, 0) + beam_search_store.predecessor_beams.index_fill_(0, seq_slots_long, 0) + beam_search_store.first_finish_reasons.index_fill_( + 0, seq_slots_long, FinishReason.NOT_FINISHED.value + ) + beam_search_store.original_tokens.index_fill_(0, seq_slots_long, 0) + beam_search_store.beam_gen_lengths.index_fill_(0, seq_slots_long, 0) + beam_search_store.prompt_lens.index_copy_(0, seq_slots_long, prompt_lens_cuda) + beam_search_store.batch_dones.index_fill_(0, seq_slots_long, False) + # The CBA tensors only exist once an exhaustive-early_stopping request has + # been admitted; nothing reads them before that, so skip the reset. + cba = beam_search_store.cba + if cba is not None: + cba.cba_tokens.index_fill_(0, seq_slots_long, BEAM_SEARCH_PAD_TOKEN) + cba.cba_cum_log_probs.index_fill_(0, seq_slots_long, 0) + cba.cba_normed_scores.index_fill_(0, seq_slots_long, float("-inf")) + cba.cba_lengths.index_fill_(0, seq_slots_long, 0) + cba.original_log_probs.index_fill_(0, seq_slots_long, 0) + cba.cba_log_probs.index_fill_(0, seq_slots_long, 0) + cba.cba_caps.index_copy_(0, seq_slots_long, beam_caps_cuda) + + +def _request_uses_cba(request: LlmRequest) -> bool: + """Whether this beam-search request runs on the candidate-beams-array + path (exhaustive early_stopping modes).""" + early_stopping = _unwrap_singleton( + cast(Optional[list[int]], request.sampling_config.early_stopping) + ) + return BeamSearchEarlyStop.from_raw(early_stopping) != BeamSearchEarlyStop.TRUE + + +def _prepare_beam_history_cba( + request: LlmRequest, + *, + cba_group: CBAGroupHost, +) -> BeamHistoryBuilder | None: + """CBA-mode variant of ``_prepare_beam_history``. + + The final beams are the top ``beam_width`` of (CBA finished paths | + current active slot paths) ranked by length-normalized score: the + unfinished active paths are inserted into the CBA and everything is + ranked together by normed score. All device state arrives through the + group-level host snapshot (``cba_group``, see _prepare_cba_group_host); + this function only slices host rows. + """ + num_tokens = request.max_beam_num_tokens + 1 # last token is not yet added + prompt_length = request.py_prompt_len + num_generated_tokens = num_tokens - prompt_length + num_beams = request.py_beam_width + + if num_generated_tokens == 0 or request.state == LlmRequestState.GENERATION_COMPLETE: + return None + + slot = request.py_seq_slot + assert slot is not None + row = cba_group.pos[slot] + # Active beams currently in the slots: the input width of the current + # step (== num_beams except for variable-beam-width requests). + active_width = _get_beam_width_in(request) + return_log_probs = request.py_return_log_probs + + length_penalty = ( + _unwrap_singleton(cast(Optional[list[float]], request.sampling_config.length_penalty)) + or 0.0 + ) + + def _builder() -> BeamHistory | None: + if not cba_group.should_stop[row].item(): + return None + + cache_indirection = cba_group.cache_indirection[ + row, :active_width, prompt_length:num_tokens + ] + current_path = cba_group.original_tokens[row, :active_width, prompt_length:num_tokens] + active_cum = cba_group.cum[row, :active_width] + cba_tokens = cba_group.cba_tokens[row] + cba_cum = cba_group.cba_cum[row] + cba_normed = cba_group.cba_normed[row] + cba_lengths = cba_group.cba_lengths[row] + + active_path = _gather_beam_path( + current_path=current_path, cache_indirection=cache_indirection + ) + active_normed = active_cum + if length_penalty != 0.0: + active_normed = active_cum / float(num_generated_tokens) ** length_penalty + active_lp_path: torch.Tensor | None = None + cba_log_probs: torch.Tensor | None = None + if return_log_probs: + assert cba_group.cba_log_probs is not None + assert cba_group.original_log_probs is not None + cba_log_probs = cba_group.cba_log_probs[row] + current_lp_path = cba_group.original_log_probs[ + row, :active_width, prompt_length:num_tokens + ] + active_lp_path = _gather_beam_path( + current_path=current_lp_path, cache_indirection=cache_indirection + ) + + pool_width = cba_normed.size(0) + all_normed = torch.cat([cba_normed, active_normed]) + order = torch.argsort(all_normed, descending=True)[:num_beams] + + width = max(num_generated_tokens, cast(int, cba_lengths.max().item())) + tokens = torch.full((num_beams, width), BEAM_SEARCH_PAD_TOKEN, dtype=torch.int32) + cum_logprobs = torch.zeros((num_beams,), dtype=torch.float32) + log_probs: torch.Tensor | None = None + if return_log_probs: + log_probs = torch.zeros((num_beams, width), dtype=torch.float32) + for out_idx, merged_idx in enumerate(order.tolist()): + if not torch.isfinite(all_normed[merged_idx]): + continue # unreachable unless fewer finite candidates than + # output beams (early termination edge); leaves a padded row + if merged_idx < pool_width: # CBA entry + entry_len = int(cba_lengths[merged_idx].item()) + tokens[out_idx, :entry_len] = cba_tokens[merged_idx, :entry_len] + cum_logprobs[out_idx] = cba_cum[merged_idx] + if log_probs is not None: + assert cba_log_probs is not None + log_probs[out_idx, :entry_len] = cba_log_probs[merged_idx, :entry_len] + else: + active_idx = merged_idx - pool_width + tokens[out_idx, :num_generated_tokens] = active_path[active_idx] + cum_logprobs[out_idx] = active_cum[active_idx] + if log_probs is not None: + assert active_lp_path is not None + log_probs[out_idx, :num_generated_tokens] = active_lp_path[active_idx] + return BeamHistory( + tokens=tokens, + # [beam, tokens, 1]: the sampled token's logprob per position, + # matching the shape contract of _convert_logprobs_tensor_to_list. + logprobs=log_probs.unsqueeze(-1) if log_probs is not None else None, + logprobs_indices=tokens.unsqueeze(-1) if return_log_probs else None, + cum_logprobs=cum_logprobs, + ) + + return _builder + + +def _postprocess_beam_logprobs( + request: LlmRequest, + *, + cache_indirection: torch.Tensor, + log_probs_host: _BeamHistoryLogProbsSlices, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Reorder per-step beam logprobs along the cache-indirection axis. + + Concatenates the freshly-sampled per-step entries onto the + request's existing host-side logprobs buffer and gathers each + beam's history through `cache_indirection`. Returns the gathered + (logprobs, logprobs_indices, cum_logprobs) triple. + """ + current_logprobs, current_logprobs_indices = get_logprobs_from_request( + request, preallocate_extra_steps=1 + ) + # concatenate the newly generated logprobs and newly + # generated tokens to the current logprobs and logprobs indices + current_logprobs[:, -1, :].copy_(log_probs_host.sampled_log_probs) + current_logprobs_indices[:, -1, :].copy_(log_probs_host.sampled_logprobs_indices) + + # Gather the correct logprobs for each beam. + new_logprobs = torch.zeros_like(current_logprobs) + new_logprobs_indices = torch.zeros_like(current_logprobs_indices) + cache_indirection_for_logprobs = cache_indirection.unsqueeze(-1).expand( + -1, -1, current_logprobs.shape[2] + ) + torch.gather( + input=current_logprobs, + dim=0, + index=cache_indirection_for_logprobs, + out=new_logprobs, + ) + torch.gather( + input=current_logprobs_indices, + dim=0, + index=cache_indirection_for_logprobs, + out=new_logprobs_indices, + ) + return new_logprobs, new_logprobs_indices, log_probs_host.cum_logprobs + + +def _finalize_beam( + request: LlmRequest, + beam_history: BeamHistory, +) -> None: + """Update the request with the corrected tokens and logprobs for each beam. + + Args: + request: The request to update + beam_history: The beam history used to update the request + """ + + beam_width = request.py_beam_width + assert beam_history.tokens.shape[0] == beam_width, ( + f"Beam_history.tokens.shape[0] should equal beam width: \ + {beam_history.tokens.shape[0]} != {beam_width}" + ) + if request.py_return_log_probs: + assert beam_history.logprobs is not None + assert beam_history.logprobs_indices is not None + assert beam_history.cum_logprobs is not None + assert beam_history.logprobs.shape[0] == beam_width, ( + f"Beam_history.logprobs.shape[0] should equal beam width: \ + {beam_history.logprobs.shape[0]} != {beam_width}" + ) + assert beam_history.logprobs_indices.shape[0] == beam_width, ( + f"Beam_history.logprobs_indices.shape[0] should equal beam width: \ + {beam_history.logprobs_indices.shape[0]} != {beam_width}" + ) + assert beam_history.cum_logprobs.shape[0] == beam_width, ( + f"Beam_history.cum_logprobs.shape[0] should equal beam width: \ + {beam_history.cum_logprobs.shape[0]} != {beam_width}" + ) + valid_tokens = (beam_history.tokens != BEAM_SEARCH_PAD_TOKEN).sum(dim=-1).tolist() + gen_token_list = [] + gen_log_probs_list = [] + for beam_idx in range(beam_width): + beam_valid_tokens = valid_tokens[beam_idx] + gen_token_list.append(beam_history.tokens[beam_idx, :beam_valid_tokens].tolist()) + if request.py_return_log_probs: + assert beam_history.logprobs_indices is not None + assert beam_history.logprobs is not None + gen_log_probs_list.append( + convert_logprobs_tensor_to_list( + beam_history.logprobs_indices[beam_idx : beam_idx + 1, :beam_valid_tokens], + beam_history.logprobs[beam_idx : beam_idx + 1, :beam_valid_tokens], + )[0] + ) + request.set_generated_tokens(gen_token_list) + if request.py_return_log_probs: + # cum_log_probs will not change when padding with end tokens. + # Therefore, we do not need to correct it + assert beam_history.cum_logprobs is not None + request.py_result.set_log_probs( + gen_log_probs_list, cum_log_probs=beam_history.cum_logprobs.tolist() + ) + + +def _check_beam_search_stop_criteria( + request: LlmRequest, + finish_reasons: torch.Tensor, +) -> torch.Tensor: + """Check if the stop criteria is met for the request. + + Returns a boolean tensor of shape (), whose value is computed asynchronously. + """ + return (finish_reasons[: request.py_beam_width] > 0).sum() == request.py_beam_width + + +class BeamSearchHandler: + """Owns the beam-search store and the host-side state the CBA path needs. + + ``TorchSampler`` holds one instance and drives it per step. The handler + keeps the pieces that several beam-search helpers share -- the + :class:`BeamSearchStore`, ``max_seq_len``, and the lagged + ``first_finish_reasons`` snapshots used by the speculative-D2H predictor -- + so they need not be threaded through every call. + """ + + def __init__( + self, + *, + store: Optional[BeamSearchStore], + log_probs_store: LogProbsStore, + new_tokens: torch.Tensor, + max_seq_len: int, + max_num_sequences: int, + use_speculative_d2h: bool, + has_multi_token_stop_words: Callable[[LlmRequest], bool], + copy_to_host: Callable[[torch.Tensor], torch.Tensor], + make_side_stream_copier: Callable[[], AbstractContextManager["_SideStreamCopier"]], + ): + self._store = store + self._log_probs_store = log_probs_store + self._new_tokens = new_tokens + self._max_seq_len = max_seq_len + self._use_speculative_d2h = use_speculative_d2h + # Bound methods of the owning sampler: the D2H copies must share its + # worker thread and side stream, so the SamplerEvent it records covers + # them. A handler-owned stream would not be awaited by that event. + self._copy_to_host = copy_to_host + self._make_side_stream_copier = make_side_stream_copier + # Lagged per-slot first_finish_reasons for the speculative predictor, + # indexed by py_seq_slot. None for unoccupied slots or before the first + # step; all-None in default mode. + self._prev_first_finish_reasons: list[torch.Tensor | None] = [None] * max_num_sequences + # Stop-word length check lives with finish-reason handling, which also + # uses it outside beam search; injected rather than duplicated here. + self._has_multi_token_stop_words = has_multi_token_stop_words + + def clear_slot(self, slot: int) -> None: + """Drop stale predictor state from a prior occupant of ``slot``.""" + self._prev_first_finish_reasons[slot] = None + + def record_first_finish_reasons(self, slot: int, reasons: torch.Tensor | None) -> None: + """Snapshot this step's first_finish_reasons for the next step's predictor.""" + self._prev_first_finish_reasons[slot] = reasons + + def predict_is_likely_finishing( + self, + request: LlmRequest, + *, + num_generated_tokens: int, + num_tokens: int, + ) -> bool: + """Predict whether this step is likely to trigger beam history finalization. + + Returns True if any of: + 1. Length budget reached (max_new_tokens or max_seq_len). + 2. Multi-token stop_words configured (forces finalization). + 3. Lagged first_finish_reasons shows any beam finished previously. + + Known miss: all beams hit end_id on the same step from a clean state. + """ + if num_generated_tokens >= request.py_max_new_tokens or num_tokens >= self._max_seq_len: + return True + if self._has_multi_token_stop_words(request): + return True + assert request.py_seq_slot is not None + prev = self._prev_first_finish_reasons[request.py_seq_slot] + # FinishReason.NOT_FINISHED == 0, so a nonzero entry implies that + # some beam has already finished. + if prev is not None and prev.any().item(): + return True + return False + + @nvtx_range("_prepare_cba_group_host") + def prepare_cba_group_host( + self, + requests: list[LlmRequest], + finish_reasons: torch.Tensor, + d2h_copier: Callable[[torch.Tensor], torch.Tensor], + ) -> Optional[CBAGroupHost]: + """Batch the per-step D2H state needed to finalize CBA-mode requests. + + The per-request variant issued ~8 small copies per request per step, + which is host-call-count bound; one batched copy per tensor for the + whole group replaces them (the builders slice the host rows). + """ + cba_requests = [request for request in requests if _request_uses_cba(request)] + if not cba_requests: + return None + store = self._store + assert store is not None + # cba_requests is non-empty here, so the tensors were allocated at + # admission (see BeamSearchHandler.ensure_cba_for_requests). + cba = store.cba + assert cba is not None, "CBA tensors must be allocated before a CBA step" + slots = [request.py_seq_slot for request in cba_requests] + assert all(slot is not None for slot in slots), "CBA requests must have seq slots" + widths = [request.py_beam_width for request in cba_requests] + both_host = torch.tensor([slots, widths], dtype=torch.int64, pin_memory=prefer_pinned()) + both_cuda = both_host.to(device="cuda", non_blocking=True) + slots_cuda, widths_cuda = both_cuda[0], both_cuda[1] + num_tokens_max = max(request.max_beam_num_tokens + 1 for request in cba_requests) + attn_width = min(store.cache_indirection.size(-1), num_tokens_max) + snap_width = min( + cba.cba_tokens.size(-1), + max( + request.max_beam_num_tokens + 2 - request.py_prompt_len for request in cba_requests + ), + ) + group_reasons = finish_reasons[slots_cuda] + should_stop = ( + (group_reasons > 0) + | ( + torch.arange(group_reasons.size(1), device=group_reasons.device).view(1, -1) + >= widths_cuda.view(-1, 1) + ) + ).all(dim=1) + return_log_probs = any(request.py_return_log_probs for request in cba_requests) + return CBAGroupHost( + pos={cast(int, slot): i for i, slot in enumerate(slots)}, + should_stop=d2h_copier(should_stop), + cache_indirection=d2h_copier(store.cache_indirection[slots_cuda, :, :attn_width]), + original_tokens=d2h_copier(store.original_tokens[slots_cuda, :, :attn_width]), + cum=d2h_copier(store.cum_log_probs[slots_cuda]), + cba_tokens=d2h_copier(cba.cba_tokens[slots_cuda, :, :snap_width]), + cba_cum=d2h_copier(cba.cba_cum_log_probs[slots_cuda]), + cba_normed=d2h_copier(cba.cba_normed_scores[slots_cuda]), + cba_lengths=d2h_copier(cba.cba_lengths[slots_cuda]), + original_log_probs=( + d2h_copier(cba.original_log_probs[slots_cuda, :, :attn_width]) + if return_log_probs + else None + ), + cba_log_probs=( + d2h_copier(cba.cba_log_probs[slots_cuda, :, :snap_width]) + if return_log_probs + else None + ), + ) + + def _prepare_beam_history( + self, + request: LlmRequest, + *, + finish_reasons: torch.Tensor, + d2h_copier: Callable[[torch.Tensor], torch.Tensor], + cba_group: Optional[CBAGroupHost] = None, + ) -> BeamHistoryBuilder | None: + """Correct the stored tokens for each beam and return it as a BeamHistory object. + + Beam Search sampling only adds new tokens to the beam. + However during beam search, a beam may change its previously sampled tokens. + This function corrects the stored tokens for each beam to match the expected tokens. + If logprobs are requested, the function also corrects the stored logprobs for each beam. + The function returns a BeamHistory object that contains the corrected tokens and logprobs for each beam. + + D2H copies are issued through `d2h_copier`. When + `_use_speculative_beam_history_d2h` is set, a host-side predictor + decides per step whether to stage copies via `d2h_copier`; + predictor misses fall back to a synchronous `.cpu()` inside + `_builder`. Otherwise, copies are issued unconditionally. + + Note: To defer the decision whether or not to skip BeamHistory construction until update_requests(), only + a builder (BeamHistoryBuilder) is returned here. The builder contains host tensors which are + being populated asynchronously. Hence, it can only be invoked after async D2H copies have completed, + e.g., after awaiting state.sampler_event in update_requests. + + arguments: + request: The request to create the beam history for + finish_reasons: The first finish reason encountered for each beam of the request. + Shape: (max_tokens, max_beam_width) + d2h_copier: Callable performing the D2H copy. + """ + if _request_uses_cba(request): + assert cba_group is not None + return _prepare_beam_history_cba(request, cba_group=cba_group) + + # Gather data used for skipping beam history processing + need_finalize_due_to_stop_words = self._has_multi_token_stop_words(request) + if need_finalize_due_to_stop_words: + need_history = torch.tensor(True) + else: + should_stop = _check_beam_search_stop_criteria( + request, + finish_reasons=finish_reasons, + ) + need_history = should_stop + # enqueue async D2H copy + need_history = self._copy_to_host(need_history) + + num_tokens = request.max_beam_num_tokens + 1 # last token is not yet added + prompt_length = request.py_prompt_len + num_generated_tokens = num_tokens - prompt_length + num_beams = request.py_beam_width + + if num_generated_tokens == 0 or request.state == LlmRequestState.GENERATION_COMPLETE: + # early return if no tokens have been generated yet or the request is already finished + return None + + beam_search_store = self._store + assert beam_search_store is not None + + log_probs_device: _BeamHistoryLogProbsSlices | None = None + if request.py_return_log_probs: + log_probs_store = self._log_probs_store + log_probs_device = _BeamHistoryLogProbsSlices( + sampled_log_probs=log_probs_store.sampled_log_probs[ + request.py_seq_slot, :num_beams + ].view(-1, 1), + sampled_logprobs_indices=self._new_tokens[0, request.py_seq_slot, :num_beams].view( + -1, 1 + ), + cum_logprobs=beam_search_store.cum_log_probs[request.py_seq_slot, :num_beams], + ) + device_slices = _BeamHistoryTensors( + cache_indirection=beam_search_store.cache_indirection[ + request.py_seq_slot, :num_beams, prompt_length:num_tokens + ], + current_path=beam_search_store.original_tokens[ + request.py_seq_slot, :num_beams, prompt_length:num_tokens + ], + log_probs=log_probs_device, + ) + + # In speculative mode, the predictor may skip the copy; otherwise + # always copy. `host_snapshot is None` triggers the .cpu() fallback + # in `_builder`, which can only happen on a predictor miss. + issue_copy = not self._use_speculative_d2h or self.predict_is_likely_finishing( + request, + num_generated_tokens=num_generated_tokens, + num_tokens=num_tokens, + ) + + host_snapshot: _BeamHistoryTensors | None = None + if issue_copy: + log_probs_host: _BeamHistoryLogProbsSlices | None = None + if device_slices.log_probs is not None: + log_probs_host = _BeamHistoryLogProbsSlices( + sampled_log_probs=d2h_copier(device_slices.log_probs.sampled_log_probs), + sampled_logprobs_indices=d2h_copier( + device_slices.log_probs.sampled_logprobs_indices + ), + cum_logprobs=d2h_copier(device_slices.log_probs.cum_logprobs), + ) + host_snapshot = _BeamHistoryTensors( + cache_indirection=d2h_copier(device_slices.cache_indirection), + current_path=d2h_copier(device_slices.current_path), + log_probs=log_probs_host, + ) + + def _builder() -> BeamHistory | None: + if not need_history.item(): + return None + + if host_snapshot is not None: + cache_indirection = host_snapshot.cache_indirection + current_path = host_snapshot.current_path + log_probs_host = host_snapshot.log_probs + else: + # Predictor-miss fallback: synchronous .cpu() on the main stream. + cache_indirection = device_slices.cache_indirection.cpu() + current_path = device_slices.current_path.cpu() + log_probs_host = None + if device_slices.log_probs is not None: + log_probs_host = _BeamHistoryLogProbsSlices( + sampled_log_probs=device_slices.log_probs.sampled_log_probs.cpu(), + sampled_logprobs_indices=( + device_slices.log_probs.sampled_logprobs_indices.cpu() + ), + cum_logprobs=device_slices.log_probs.cum_logprobs.cpu(), + ) + + new_path = _gather_beam_path( + current_path=current_path, cache_indirection=cache_indirection + ) + new_logprobs: torch.Tensor | None = None + new_logprobs_indices: torch.Tensor | None = None + cum_logprobs_out: torch.Tensor | None = None + if log_probs_host is not None: + new_logprobs, new_logprobs_indices, cum_logprobs_out = _postprocess_beam_logprobs( + request, + cache_indirection=cache_indirection, + log_probs_host=log_probs_host, + ) + + return BeamHistory( + tokens=new_path, + logprobs=new_logprobs, + logprobs_indices=new_logprobs_indices, + cum_logprobs=cum_logprobs_out, + ) + + return _builder + + def prepare_beam_histories( + self, + requests: list[LlmRequest], + finish_reasons: torch.Tensor, + ) -> tuple[list[BeamHistoryBuilder | None], torch.cuda.Event | None]: + """Create the corrected tokens and logprobs for each beam of a request. + + The builders returned by this function create a beam history object containing + the corrected tokens and logprobs for each beam of a request. + + Returns (builders, side_stream_event). side_stream_event is set + only when the speculative path queued copies; the caller must + forward it to _record_sampler_event so SamplerEvent.synchronize + awaits the side stream before any builder is invoked. + """ + # Single `with` for both modes; nullcontext yields None. + copier_ctx: AbstractContextManager["_SideStreamCopier | None"] = ( + self._make_side_stream_copier() if self._use_speculative_d2h else nullcontext() + ) + with copier_ctx as copier: + d2h_copier: Callable[[torch.Tensor], torch.Tensor] = ( + copier.stage_copy_to_host if copier is not None else self._copy_to_host + ) + cba_group = self.prepare_cba_group_host(requests, finish_reasons, d2h_copier) + builders = [ + self._prepare_beam_history( + req, + finish_reasons=finish_reasons[req.py_seq_slot], + d2h_copier=d2h_copier, + cba_group=cba_group, + ) + for req in requests + ] + side_stream_event = copier.event if copier is not None else None + return builders, side_stream_event diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py index 6ee643336eb7..af265b33f987 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py @@ -64,6 +64,27 @@ def __getattr__(self, name: str) -> Any: SeedOrTensor = Union[int, torch.Tensor] +@_compiler_disable +def radix_topk_op( + values: torch.Tensor, + k: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Sorted top-k via flashinfer's radix-select kernel. + + Drop-in for ``torch.topk(values, k, dim=-1, sorted=True)`` on 2D inputs: + returns ``(values, indices)`` with values descending and int64 indices; + like torch.topk, the index order among equal values is unspecified. + O(n) radix select — much faster than torch.topk on large rows, but with a + fixed per-call cost that torch.topk undercuts on small rows (see + ``beam_search._beam_topk`` for the size-dispatching entry point). + + ``deterministic=True``: the default collect pass races, which makes the + order among equal values vary run to run. This op sits on the default beam + path, where reproducible tie order matters for triage. + """ + return flashinfer.top_k(values, k, sorted=True, deterministic=True) + + @_compiler_disable def top_k_top_p_sampling_from_logits_op( logits: torch.Tensor, diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py index e1811c6a4640..88b682813c8d 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/ops/vanilla.py @@ -25,9 +25,6 @@ import torch from tensorrt_llm._utils import prefer_pinned -from tensorrt_llm.bindings.executor import FinishReason - -BEAM_SEARCH_PAD_TOKEN = -1 @dataclass(kw_only=True) @@ -35,22 +32,6 @@ class StrategyMetadata: """Base class for per-strategy-group metadata passed into sample().""" -@dataclass(kw_only=True) -class BeamSearchMetadata(StrategyMetadata): - """Stateful tensors required by beam_search_sampling_batch.""" - - cache_indirection: torch.Tensor - cache_indirection_buffer: torch.Tensor - cum_log_probs: torch.Tensor - new_log_probs: torch.Tensor - seq_slots: torch.Tensor - seq_lens: torch.Tensor - finished_beams: torch.Tensor - predecessor_beams: torch.Tensor - seq_offsets: torch.Tensor - beam_idx_arange: torch.Tensor - - def min_p_renorm_probs( probs: torch.Tensor, min_p: torch.Tensor | float, @@ -163,119 +144,6 @@ def greedy_search_sampling_batch( return next_tokens, softmax -def _update_cache_indirection_buffer( - cache_indirection_input: torch.Tensor, - cache_indirection_output: torch.Tensor, - seq_slots: torch.Tensor, -) -> None: - assert cache_indirection_input.device == cache_indirection_output.device - cache_indirection_input.index_copy_(0, seq_slots, cache_indirection_output[seq_slots]) - - -def beam_search_sampling_batch( - logits: torch.Tensor, - *, - beam_width_in: int, - beam_width_out: int, - beam_search_args: BeamSearchMetadata, - temperature: float | None, - return_probs: bool = True, -) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - """Sample beam_width tokens for each request in parallel.""" - logits_dim = logits.dim() - assert logits_dim == 2, "logits should be 2D: [batch_size * beam_width, vocab_size]" - batch_size, vocab_size = logits.size() - batch_size = batch_size // beam_width_in - - logits = logits.view(batch_size, beam_width_in, vocab_size) - if temperature is not None and temperature != 0: - logits = logits / max(temperature, 1e-5) - softmax: Optional[torch.Tensor] = None - if return_probs: - softmax = torch.softmax(logits, dim=-1) - _update_cache_indirection_buffer( - beam_search_args.cache_indirection_buffer, - beam_search_args.cache_indirection, - beam_search_args.seq_slots, - ) - assert batch_size == beam_search_args.seq_slots.size(0) - - logprobs = torch.log_softmax(logits, dim=-1) - - finished_beams_mask = ( - beam_search_args.finished_beams[beam_search_args.seq_slots, :beam_width_in] - != FinishReason.NOT_FINISHED.value - ) - finished_beams_mask_expanded = finished_beams_mask.unsqueeze(-1).expand( - -1, -1, logprobs.size(-1) - ) - logprobs = torch.where(finished_beams_mask_expanded, float("-inf"), logprobs) - logprobs[..., 0] = torch.where(finished_beams_mask, 0, logprobs[..., 0]) - - logprobs += beam_search_args.cum_log_probs.unsqueeze(-1)[ - beam_search_args.seq_slots, :beam_width_in - ] - - logprobs = logprobs.view(batch_size, beam_width_in * vocab_size) - sorted_logprobs, sorted_indices = torch.topk(logprobs, k=beam_width_out, sorted=True, dim=-1) - - next_tokens = sorted_indices.to(torch.int32) - - predecessor_beam = next_tokens // vocab_size - beam_search_args.predecessor_beams[beam_search_args.seq_slots, :beam_width_out] = ( - predecessor_beam - ) - - max_beam_width = beam_search_args.finished_beams.size(1) - finished_beams = beam_search_args.finished_beams[beam_search_args.seq_slots].view(-1) - - offset_predecessor_beam = predecessor_beam + beam_search_args.seq_offsets[ - : predecessor_beam.size(0) - ].unsqueeze(1) - finished_beams = finished_beams[offset_predecessor_beam] - beam_search_args.finished_beams[beam_search_args.seq_slots] = finished_beams.view( - batch_size, max_beam_width - ) - - cache_indirection = beam_search_args.cache_indirection[ - beam_search_args.seq_slots, :beam_width_out - ] - cache_indirection_buffer = beam_search_args.cache_indirection_buffer[ - beam_search_args.seq_slots, :beam_width_in - ] - torch.gather( - cache_indirection_buffer, - dim=1, - index=predecessor_beam.unsqueeze(2).expand(-1, -1, cache_indirection.size(2)), - out=cache_indirection, - ) - - index = beam_search_args.seq_lens.view(-1, 1, 1).expand(-1, beam_width_out, 1) - src = ( - beam_search_args.beam_idx_arange[:beam_width_out] - .view(1, beam_width_out, 1) - .expand(batch_size, beam_width_out, 1) - ) - cache_indirection.scatter_(2, index, src) - - beam_search_args.cache_indirection[beam_search_args.seq_slots, :beam_width_out] = ( - cache_indirection - ) - - next_tokens = next_tokens % vocab_size - ended_predecessor_mask = torch.gather(dim=1, index=predecessor_beam, input=finished_beams_mask) - next_tokens = torch.where(ended_predecessor_mask, BEAM_SEARCH_PAD_TOKEN, next_tokens) - - old_cum_log_probs = beam_search_args.cum_log_probs[beam_search_args.seq_slots].view(-1) - beam_search_args.new_log_probs[beam_search_args.seq_slots, :beam_width_out] = ( - sorted_logprobs[:, :beam_width_out] - old_cum_log_probs[offset_predecessor_beam] - ) - beam_search_args.cum_log_probs[beam_search_args.seq_slots, :beam_width_out] = sorted_logprobs[ - :, :beam_width_out - ] - return next_tokens, softmax - - def get_rejected_indices( draft_probs: torch.Tensor, target_probs: torch.Tensor, diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index ed58c84ac854..2617fa4bfab6 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -18,7 +18,7 @@ from collections import defaultdict from collections.abc import Iterable, Iterator from concurrent import futures -from contextlib import AbstractContextManager, contextmanager, nullcontext +from contextlib import contextmanager from dataclasses import dataclass, field from itertools import repeat from typing import ( @@ -87,13 +87,18 @@ from ..llm_request import LlmRequest, LlmRequestState, get_draft_token_length from ..resource_manager import ResourceManager, ResourceManagerType from ..scheduler import ScheduledRequests +from .beam_search import ( + BeamHistoryBuilder, + BeamSearchHandler, + _finalize_beam, + _prepare_beam_search, + _request_uses_cba, +) from .finish_reasons import FinishReasonsHandler from .logprobs import ( LogProbsState, LogProbsStateList, LogProbsStore, - convert_logprobs_tensor_to_list, - get_logprobs_from_request, store_logprobs_list_to_request, ) from .penalties import PenaltyHandler, has_occurrence_penalty @@ -108,9 +113,11 @@ request_random_seed, ) from .sampler_strategy import ( - BEAM_SEARCH_PAD_TOKEN, GREEDY, + BeamHistory, BeamSearchMetadata, + BeamSearchStore, + CBAState, FlashInferGroupedStrategySampler, Fusions, GenericStrategyKeyType, @@ -156,6 +163,12 @@ class SampleStateTensors: log_probs: torch.Tensor | None = None +GenericSampleStateTensorsHost = TypeVar("GenericSampleStateTensorsHost", bound=SampleStateTensors) +GenericSampleStateTensorsDevice = TypeVar( + "GenericSampleStateTensorsDevice", bound=SampleStateTensors +) + + @dataclass(kw_only=True) class SamplerEvent: cuda_event: torch.cuda.Event @@ -171,12 +184,6 @@ def synchronize(self) -> None: self.side_stream_event.synchronize() -GenericSampleStateTensorsHost = TypeVar("GenericSampleStateTensorsHost", bound=SampleStateTensors) -GenericSampleStateTensorsDevice = TypeVar( - "GenericSampleStateTensorsDevice", bound=SampleStateTensors -) - - @dataclass(kw_only=True) class SampleState(Generic[GenericSampleStateTensorsHost, GenericSampleStateTensorsDevice]): requests: list[LlmRequest] @@ -1148,27 +1155,6 @@ def __init__( raise ValueError(f"Invalid dim_order: {dim_order}") -@dataclass(kw_only=True) -class BeamHistory: - """ - Beam history class for beam search. - This class is used to store the corrected tokens and logprobs for each beam. - It is used to update the beam history for each beam. - """ - - tokens: torch.Tensor - logprobs: torch.Tensor | None = None - logprobs_indices: torch.Tensor | None = None - cum_logprobs: torch.Tensor | None = None - - -BeamHistoryBuilder: TypeAlias = Callable[[], BeamHistory | None] -"""Builder for BeamHistory. - -Used to defer possibly unnecessary host-tensor construction until update_requests(). -""" - - @dataclass(kw_only=True) class SamplingRequestsMetadata: """Metadata for the sampling requests.""" @@ -1209,38 +1195,6 @@ class SampleStateTorch(SampleState[SampleStateTensorsHostTorch, SampleStateTenso single_step_greedy: bool = False -@dataclass(kw_only=True, frozen=True) -class _BeamHistoryLogProbsSlices: - """Correlated beam-history log-prob tensors; all three fields are bound together.""" - - sampled_log_probs: torch.Tensor - sampled_logprobs_indices: torch.Tensor - cum_logprobs: torch.Tensor - - -@dataclass(kw_only=True, frozen=True) -class _BeamHistoryTensors: - """Beam-history tensor slices. - - Used to carry both device-side views (before D2H) and host-side - snapshots (after D2H). `log_probs` is bound iff log-probs are - requested. - """ - - cache_indirection: torch.Tensor - current_path: torch.Tensor - log_probs: _BeamHistoryLogProbsSlices | None - - -def _gather_beam_path( - *, current_path: torch.Tensor, cache_indirection: torch.Tensor -) -> torch.Tensor: - """Gather the correct tokens for each beam from current_path.""" - new_path = torch.zeros_like(current_path) - torch.gather(input=current_path, dim=0, index=cache_indirection, out=new_path) - return new_path - - class _SideStreamCopier: """Batch non-blocking D2H copies onto a private side stream. @@ -1442,37 +1396,6 @@ def is_generation_model(self) -> bool: return True @dataclass(kw_only=True) - class BeamSearchStore: - """Auxiliary data structures required for beam search.""" - - cache_indirection: torch.Tensor - """Shape: batch_size, beam_width, attention_size - Usage: Stores the cache indirection necessary for beam search sampling""" - cache_indirection_buffer: torch.Tensor - """Shape: batch_size, beam_width, attention_size - Usage: A second buffer used to update the cache indirection during sampling""" - cum_log_probs: torch.Tensor - """Shape: batch_size, beam_width - Usage: Stores the current cumulative logprob of each active beam for faster sampling""" - first_finish_reasons: torch.Tensor - """Shape: batch_size, beam_width - Usage: Stores the first finish reason for each beam""" - predecessor_beams: torch.Tensor - """Shape: batch_size, beam_width - Usage: Stores the predecessor beams for each beam used for stop word detection""" - original_tokens: torch.Tensor - """Shape: batch_size, beam_width, sequence_length - Usage: Stores the original tokens for each beam. - This is used to recover the original tokens for each beam when streaming is enabled""" - seq_offsets: torch.Tensor - """Shape: (max_num_sequences,), dtype int64 - Usage: Cached `arange(max_num_sequences) * max_beam_width` used by - ``beam_search_sampling_batch`` to flatten (batch_idx, beam_idx) pairs.""" - beam_idx_arange: torch.Tensor - """Shape: (max_beam_width,), dtype int32 - Usage: Cached `arange(max_beam_width)` used as the scatter source in the - per-step ``cache_indirection.scatter_``.""" - @dataclass(kw_only=True) class Store: new_tokens: torch.Tensor @@ -1480,7 +1403,7 @@ class Store: Shape: See cpp DecoderState.getAllNewTokens(). """ - beam_search_store: "TorchSampler.BeamSearchStore | None" = None + beam_search_store: "BeamSearchStore | None" = None """Holds data related to beam search.""" log_probs_store: LogProbsStore """Holds data related to log-probs handling.""" @@ -1509,30 +1432,10 @@ def _create_store(self) -> Store: beam_search_store = None if self._use_beam_search: - cache_indirection = torch.empty( - self.CACHE_INDIRECTION_SHAPE, device="cuda", dtype=torch.int - ) - cache_indirection_buffer = int_tensor(self.CACHE_INDIRECTION_SHAPE) - cum_log_probs = torch.empty( - self.CACHE_INDIRECTION_SHAPE[:-1], device="cuda", dtype=torch.float32 - ) - predecessor_beams = int_tensor(self.CACHE_INDIRECTION_SHAPE[:-1]) - original_tokens = int_tensor(self.CACHE_INDIRECTION_SHAPE) - first_finish_reasons = int_tensor(self.CACHE_INDIRECTION_SHAPE[:-1]) - seq_offsets = ( - torch.arange(self.max_num_sequences, device="cuda", dtype=torch.int64) - * self.max_beam_width - ) - beam_idx_arange = torch.arange(self.max_beam_width, device="cuda", dtype=torch.int32) - beam_search_store = self.BeamSearchStore( - cache_indirection=cache_indirection, - cache_indirection_buffer=cache_indirection_buffer, - cum_log_probs=cum_log_probs, - predecessor_beams=predecessor_beams, - original_tokens=original_tokens, - first_finish_reasons=first_finish_reasons, - seq_offsets=seq_offsets, - beam_idx_arange=beam_idx_arange, + beam_search_store = BeamSearchStore.create( + cache_indirection_shape=self.CACHE_INDIRECTION_SHAPE, + max_num_sequences=self.max_num_sequences, + max_beam_width=self.max_beam_width, ) return self.Store( new_tokens=new_tokens, @@ -1663,17 +1566,25 @@ def __init__(self, args: Args): ) self._use_speculative_beam_history_d2h = False - # 1-step-lagged host mirror of first_finish_reasons used by the - # speculative predictor, indexed by py_seq_slot. None for unoccupied - # slots or before the first step; all-None in default mode. - self._prev_first_finish_reasons_host: list[torch.Tensor | None] = [ - None - ] * self.max_num_sequences self._stable_greedy_request_ids: list[int] = [] self._stable_greedy_seq_slots: list[int] = [] self._stable_greedy_seq_slots_host: Optional[torch.Tensor] = None self._stable_greedy_seq_slots_cuda: Optional[torch.Tensor] = None + # BeamSearchHandler owns the lagged first_finish_reasons snapshots the + # speculative predictor reads, so no separate host mirror is kept here. + self._beam_search = BeamSearchHandler( + store=self.store.beam_search_store, + log_probs_store=self.store.log_probs_store, + new_tokens=self.store.new_tokens, + max_seq_len=self.max_seq_len, + max_num_sequences=self.max_num_sequences, + use_speculative_d2h=self._use_speculative_beam_history_d2h, + has_multi_token_stop_words=self._check_stop_words_length, + copy_to_host=self._copy_to_host, + make_side_stream_copier=self._make_side_stream_copier, + ) + @staticmethod def _is_draft_batch(requests: list[LlmRequest]) -> bool: """Whether this batch belongs to the draft model. @@ -2078,6 +1989,10 @@ def validate_request(self, request: LlmRequest) -> None: raise ValueError( "Beam search only supports returning the sampled logprob per token" ) + # early_stopping == TRUE (default) is served by the frozen-slot path + # (stopping once all beam slots hold finished beams is equivalent); + # the exhaustive modes by the candidate-beams-array path (see + # beam_search_sampling_batch_cba). @override @nvtx_range("setup_sampler_step") @@ -2113,15 +2028,17 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: max_prompt_len = max(max_prompt_len, request.py_prompt_len) if self._use_speculative_beam_history_d2h: # Drop stale predictor state from any prior occupant of this slot. - self._prev_first_finish_reasons_host[slot] = None + self._beam_search.clear_slot(slot) self._request_grouper.prepare_for_new_request(request, slot) self._penalty_handler.prepare_for_new_request(request, slot) max_lens = self._finish_reasons_handler.new_max_lens end_ids = self._finish_reasons_handler.new_end_ids + prompt_lens = [request.py_prompt_len for request in new_requests] + beam_caps = [request.py_beam_width for request in new_requests] # Perform updates to the stores - full_list = [seq_slots, max_lens, end_ids] + full_list = [seq_slots, max_lens, end_ids, prompt_lens, beam_caps] # perform only a single copy full_list_tensor_host = torch.tensor( full_list, device="cpu", dtype=torch.int32, pin_memory=prefer_pinned() @@ -2131,6 +2048,8 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: seq_slots_tensor_cuda = full_list_tensor_cuda[0] max_lens_tensor_cuda = full_list_tensor_cuda[1] end_ids_tensor_cuda = full_list_tensor_cuda[2] + prompt_lens_tensor_cuda = full_list_tensor_cuda[3] + beam_caps_tensor_cuda = full_list_tensor_cuda[4] # Cast to int64 once for downstream ``index_copy_`` / ``index_fill_`` calls. seq_slots_tensor_cuda_long = seq_slots_tensor_cuda.long() @@ -2154,39 +2073,19 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: if self._use_beam_search: beam_search_store = self.store.beam_search_store assert beam_search_store is not None - self._prepare_beam_search( + # Allocate the CBA tensors on the first exhaustive-early_stopping + # request; beam search with the default mode never needs them. + if any(_request_uses_cba(request) for request in new_requests): + beam_search_store.ensure_cba() + _prepare_beam_search( beam_search_store, self.store.log_probs_store, seq_slots_long=seq_slots_tensor_cuda_long, max_prompt_len=max_prompt_len, + prompt_lens_cuda=prompt_lens_tensor_cuda, + beam_caps_cuda=beam_caps_tensor_cuda, ) - @staticmethod - def _prepare_beam_search( - beam_search_store: BeamSearchStore, - log_probs_store: LogProbsStore, - seq_slots_long: torch.Tensor, - max_prompt_len: int, - ) -> None: - """Prepare the beam search buffers for the requests - - If the last context chunk is being processed, - initialize/reset the buffers for the request. - - ``seq_slots_long`` must be int64 (required by ``index_fill_``). - """ - beam_search_store.cache_indirection.narrow(2, 0, max_prompt_len).index_fill_( - 0, seq_slots_long, 0 - ) - beam_search_store.cum_log_probs.index_fill_(0, seq_slots_long, 0) - log_probs_store.sampled_log_probs.index_fill_(0, seq_slots_long, 0) - log_probs_store.sampled_log_prob_ranks.index_fill_(0, seq_slots_long, 0) - beam_search_store.predecessor_beams.index_fill_(0, seq_slots_long, 0) - beam_search_store.first_finish_reasons.index_fill_( - 0, seq_slots_long, FinishReason.NOT_FINISHED.value - ) - beam_search_store.original_tokens.index_fill_(0, seq_slots_long, 0) - @torch.inference_mode() def _process_draft_tokens_rejection_sampling( self, @@ -2318,260 +2217,6 @@ def process_draft_tokens( request, new_tokens_list=new_tokens_list, new_tokens_tensor=new_tokens_tensor ) - def _prepare_beam_history( - self, - request: LlmRequest, - *, - finish_reasons: torch.Tensor, - d2h_copier: Callable[[torch.Tensor], torch.Tensor], - ) -> BeamHistoryBuilder | None: - """Correct the stored tokens for each beam and return it as a BeamHistory object. - - Beam Search sampling only adds new tokens to the beam. - However during beam search, a beam may change its previously sampled tokens. - This function corrects the stored tokens for each beam to match the expected tokens. - If logprobs are requested, the function also corrects the stored logprobs for each beam. - The function returns a BeamHistory object that contains the corrected tokens and logprobs for each beam. - - D2H copies are issued through `d2h_copier`. When - `_use_speculative_beam_history_d2h` is set, a host-side predictor - decides per step whether to stage copies via `d2h_copier`; - predictor misses fall back to a synchronous `.cpu()` inside - `_builder`. Otherwise, copies are issued unconditionally. - - Note: To defer the decision whether or not to skip BeamHistory construction until update_requests(), only - a builder (BeamHistoryBuilder) is returned here. The builder contains host tensors which are - being populated asynchronously. Hence, it can only be invoked after async D2H copies have completed, - e.g., after awaiting state.sampler_event in update_requests. - - arguments: - request: The request to create the beam history for - finish_reasons: The first finish reason encountered for each beam of the request. - Shape: (max_tokens, max_beam_width) - d2h_copier: Callable performing the D2H copy. - """ - - # Gather data used for skipping beam history processing - need_finalize_due_to_stop_words = self._check_stop_words_length(request) - if need_finalize_due_to_stop_words: - need_history = torch.tensor(True) - else: - should_stop = self._check_beam_search_stop_criteria( - request, - finish_reasons=finish_reasons, - ) - need_history = should_stop - # enqueue async D2H copy - need_history = self._copy_to_host(need_history) - - num_tokens = request.max_beam_num_tokens + 1 # last token is not yet added - prompt_length = request.py_prompt_len - num_generated_tokens = num_tokens - prompt_length - num_beams = request.py_beam_width - - if num_generated_tokens == 0 or request.state == LlmRequestState.GENERATION_COMPLETE: - # early return if no tokens have been generated yet or the request is already finished - return None - - beam_search_store = self.store.beam_search_store - assert beam_search_store is not None - - log_probs_device: _BeamHistoryLogProbsSlices | None = None - if request.py_return_log_probs: - log_probs_store = self.store.log_probs_store - log_probs_device = _BeamHistoryLogProbsSlices( - sampled_log_probs=log_probs_store.sampled_log_probs[ - request.py_seq_slot, :num_beams - ].view(-1, 1), - sampled_logprobs_indices=self.store.new_tokens[ - 0, request.py_seq_slot, :num_beams - ].view(-1, 1), - cum_logprobs=beam_search_store.cum_log_probs[request.py_seq_slot, :num_beams], - ) - device_slices = _BeamHistoryTensors( - cache_indirection=beam_search_store.cache_indirection[ - request.py_seq_slot, :num_beams, prompt_length:num_tokens - ], - current_path=beam_search_store.original_tokens[ - request.py_seq_slot, :num_beams, prompt_length:num_tokens - ], - log_probs=log_probs_device, - ) - - # In speculative mode, the predictor may skip the copy; otherwise - # always copy. `host_snapshot is None` triggers the .cpu() fallback - # in `_builder`, which can only happen on a predictor miss. - issue_copy = ( - not self._use_speculative_beam_history_d2h - or self._predict_beam_search_is_likely_finishing( - request, - num_generated_tokens=num_generated_tokens, - num_tokens=num_tokens, - ) - ) - - host_snapshot: _BeamHistoryTensors | None = None - if issue_copy: - log_probs_host: _BeamHistoryLogProbsSlices | None = None - if device_slices.log_probs is not None: - log_probs_host = _BeamHistoryLogProbsSlices( - sampled_log_probs=d2h_copier(device_slices.log_probs.sampled_log_probs), - sampled_logprobs_indices=d2h_copier( - device_slices.log_probs.sampled_logprobs_indices - ), - cum_logprobs=d2h_copier(device_slices.log_probs.cum_logprobs), - ) - host_snapshot = _BeamHistoryTensors( - cache_indirection=d2h_copier(device_slices.cache_indirection), - current_path=d2h_copier(device_slices.current_path), - log_probs=log_probs_host, - ) - - def _builder() -> BeamHistory | None: - if not need_history.item(): - return None - - if host_snapshot is not None: - cache_indirection = host_snapshot.cache_indirection - current_path = host_snapshot.current_path - log_probs_host = host_snapshot.log_probs - else: - # Predictor-miss fallback: synchronous .cpu() on the main stream. - cache_indirection = device_slices.cache_indirection.cpu() - current_path = device_slices.current_path.cpu() - log_probs_host = None - if device_slices.log_probs is not None: - log_probs_host = _BeamHistoryLogProbsSlices( - sampled_log_probs=device_slices.log_probs.sampled_log_probs.cpu(), - sampled_logprobs_indices=( - device_slices.log_probs.sampled_logprobs_indices.cpu() - ), - cum_logprobs=device_slices.log_probs.cum_logprobs.cpu(), - ) - - new_path = _gather_beam_path( - current_path=current_path, cache_indirection=cache_indirection - ) - new_logprobs: torch.Tensor | None = None - new_logprobs_indices: torch.Tensor | None = None - cum_logprobs_out: torch.Tensor | None = None - if log_probs_host is not None: - new_logprobs, new_logprobs_indices, cum_logprobs_out = ( - self._postprocess_beam_logprobs( - request, - cache_indirection=cache_indirection, - log_probs_host=log_probs_host, - ) - ) - - return BeamHistory( - tokens=new_path, - logprobs=new_logprobs, - logprobs_indices=new_logprobs_indices, - cum_logprobs=cum_logprobs_out, - ) - - return _builder - - def _postprocess_beam_logprobs( - self, - request: LlmRequest, - *, - cache_indirection: torch.Tensor, - log_probs_host: _BeamHistoryLogProbsSlices, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Reorder per-step beam logprobs along the cache-indirection axis. - - Concatenates the freshly-sampled per-step entries onto the - request's existing host-side logprobs buffer and gathers each - beam's history through `cache_indirection`. Returns the gathered - (logprobs, logprobs_indices, cum_logprobs) triple. - """ - current_logprobs, current_logprobs_indices = get_logprobs_from_request( - request, preallocate_extra_steps=1 - ) - # concatenate the newly generated logprobs and newly - # generated tokens to the current logprobs and logprobs indices - current_logprobs[:, -1, :].copy_(log_probs_host.sampled_log_probs) - current_logprobs_indices[:, -1, :].copy_(log_probs_host.sampled_logprobs_indices) - - # Gather the correct logprobs for each beam. - new_logprobs = torch.zeros_like(current_logprobs) - new_logprobs_indices = torch.zeros_like(current_logprobs_indices) - cache_indirection_for_logprobs = cache_indirection.unsqueeze(-1).expand( - -1, -1, current_logprobs.shape[2] - ) - torch.gather( - input=current_logprobs, - dim=0, - index=cache_indirection_for_logprobs, - out=new_logprobs, - ) - torch.gather( - input=current_logprobs_indices, - dim=0, - index=cache_indirection_for_logprobs, - out=new_logprobs_indices, - ) - return new_logprobs, new_logprobs_indices, log_probs_host.cum_logprobs - - def _finalize_beam( - self, - request: LlmRequest, - beam_history: BeamHistory, - ) -> None: - """Update the request with the corrected tokens and logprobs for each beam. - - Args: - request: The request to update - beam_history: The beam history used to update the request - """ - - beam_width = request.py_beam_width - assert beam_history.tokens.shape[0] == beam_width, ( - f"Beam_history.tokens.shape[0] should equal beam width: \ - {beam_history.tokens.shape[0]} != {beam_width}" - ) - if request.py_return_log_probs: - assert beam_history.logprobs is not None - assert beam_history.logprobs_indices is not None - assert beam_history.cum_logprobs is not None - assert beam_history.logprobs.shape[0] == beam_width, ( - f"Beam_history.logprobs.shape[0] should equal beam width: \ - {beam_history.logprobs.shape[0]} != {beam_width}" - ) - assert beam_history.logprobs_indices.shape[0] == beam_width, ( - f"Beam_history.logprobs_indices.shape[0] should equal beam width: \ - {beam_history.logprobs_indices.shape[0]} != {beam_width}" - ) - assert beam_history.cum_logprobs.shape[0] == beam_width, ( - f"Beam_history.cum_logprobs.shape[0] should equal beam width: \ - {beam_history.cum_logprobs.shape[0]} != {beam_width}" - ) - valid_tokens = (beam_history.tokens != BEAM_SEARCH_PAD_TOKEN).sum(dim=-1).tolist() - gen_token_list = [] - gen_log_probs_list = [] - for beam_idx in range(beam_width): - beam_valid_tokens = valid_tokens[beam_idx] - gen_token_list.append(beam_history.tokens[beam_idx, :beam_valid_tokens].tolist()) - if request.py_return_log_probs: - assert beam_history.logprobs_indices is not None - assert beam_history.logprobs is not None - gen_log_probs_list.append( - convert_logprobs_tensor_to_list( - beam_history.logprobs_indices[beam_idx : beam_idx + 1, :beam_valid_tokens], - beam_history.logprobs[beam_idx : beam_idx + 1, :beam_valid_tokens], - )[0] - ) - request.set_generated_tokens(gen_token_list) - if request.py_return_log_probs: - # cum_log_probs will not change when padding with end tokens. - # Therefore, we do not need to correct it - assert beam_history.cum_logprobs is not None - request.py_result.set_log_probs( - gen_log_probs_list, cum_log_probs=beam_history.cum_logprobs.tolist() - ) - def _add_metadata_to_grouped_requests( self, requests: list[LlmRequest], @@ -2622,6 +2267,33 @@ def _add_metadata_to_grouped_requests( predecessor_beams=beam_search_store.predecessor_beams, seq_offsets=beam_search_store.seq_offsets, beam_idx_arange=beam_search_store.beam_idx_arange, + beam_gen_lengths=beam_search_store.beam_gen_lengths, + stop_past_tokens=self._finish_reasons_handler.store.past_tokens_cuda, + # None unless an exhaustive-early_stopping request has been + # admitted; the CBA tensors are not allocated before that. + cba=None + if beam_search_store.cba is None + else CBAState( + end_ids=self._finish_reasons_handler.store.end_ids_cuda, + prompt_lens=beam_search_store.prompt_lens, + original_tokens=beam_search_store.original_tokens, + batch_dones=beam_search_store.batch_dones, + cba_tokens=beam_search_store.cba.cba_tokens, + cba_cum_log_probs=beam_search_store.cba.cba_cum_log_probs, + cba_normed_scores=beam_search_store.cba.cba_normed_scores, + cba_lengths=beam_search_store.cba.cba_lengths, + original_log_probs=beam_search_store.cba.original_log_probs, + cba_log_probs=beam_search_store.cba.cba_log_probs, + cba_caps=beam_search_store.cba.cba_caps, + max_seq_len=self.max_seq_len, + max_gen_len=max( + ( + requests[i].max_beam_num_tokens + 2 - requests[i].py_prompt_len + for i in value.indices.tolist() + ), + default=0, + ), + ), ) elif metadata_type is TopPDecayMetadata: metadata = self._top_p_decay.build_metadata( @@ -2643,17 +2315,6 @@ def _add_metadata_to_grouped_requests( ) return grouped_requests_with_metadata - @staticmethod - def _check_beam_search_stop_criteria( - request: LlmRequest, - finish_reasons: torch.Tensor, - ) -> torch.Tensor: - """Check if the stop criteria is met for the request. - - Returns a boolean tensor of shape (), whose value is computed asynchronously. - """ - return (finish_reasons[: request.py_beam_width] > 0).sum() == request.py_beam_width - @staticmethod def _check_stop_words_length(request: LlmRequest) -> bool: """Check if the stop words length is greater than 1""" @@ -2670,71 +2331,6 @@ def _check_stop_words_length(request: LlmRequest) -> bool: return longest_stop_word_len > 1 return False - def _predict_beam_search_is_likely_finishing( - self, - request: LlmRequest, - *, - num_generated_tokens: int, - num_tokens: int, - ) -> bool: - """Predict whether this step is likely to trigger beam history finalization. - - Returns True if any of: - 1. Length budget reached (max_new_tokens or max_seq_len). - 2. Multi-token stop_words configured (forces finalization). - 3. Lagged first_finish_reasons shows any beam finished previously. - - Known miss: all beams hit end_id on the same step from a clean state. - """ - if num_generated_tokens >= request.py_max_new_tokens or num_tokens >= self.max_seq_len: - return True - if self._check_stop_words_length(request): - return True - assert request.py_seq_slot is not None - prev = self._prev_first_finish_reasons_host[request.py_seq_slot] - # FinishReason.NOT_FINISHED == 0, so a nonzero entry implies that - # some beam has already finished. - if prev is not None and prev.any().item(): - return True - return False - - @nvtx_range("maybe_create_beam_histories") - def _prepare_beam_histories( - self, - requests: list[LlmRequest], - finish_reasons: torch.Tensor, - ) -> tuple[list[BeamHistoryBuilder | None], torch.cuda.Event | None]: - """Create the corrected tokens and logprobs for each beam of a request. - - The builders returned by this function create a beam history object containing - the corrected tokens and logprobs for each beam of a request. - - Returns (builders, side_stream_event). side_stream_event is set - only when the speculative path queued copies; the caller must - forward it to _record_sampler_event so SamplerEvent.synchronize - awaits the side stream before any builder is invoked. - """ - # Single `with` for both modes; nullcontext yields None. - copier_ctx: AbstractContextManager[_SideStreamCopier | None] = ( - self._make_side_stream_copier() - if self._use_speculative_beam_history_d2h - else nullcontext() - ) - with copier_ctx as copier: - d2h_copier: Callable[[torch.Tensor], torch.Tensor] = ( - copier.stage_copy_to_host if copier is not None else self._copy_to_host - ) - builders = [ - self._prepare_beam_history( - req, - finish_reasons=finish_reasons[req.py_seq_slot], - d2h_copier=d2h_copier, - ) - for req in requests - ] - side_stream_event = copier.event if copier is not None else None - return builders, side_stream_event - @override @nvtx_range("update_requests") @torch.inference_mode() @@ -2839,7 +2435,7 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None: if req.py_beam_width > 1: if (beam_history := _maybe_build_beam_history(req_idx)) is not None: - self._finalize_beam(req, beam_history) + _finalize_beam(req, beam_history) else: for beam_idx in range(req.py_beam_width): # Beam search does not support speculative decoding. @@ -2853,8 +2449,8 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None: if self._use_speculative_beam_history_d2h: # Snapshot for the next step's predictor. assert req.py_seq_slot is not None - self._prev_first_finish_reasons_host[req.py_seq_slot] = ( - first_finish_reasons_host[req.py_seq_slot] + self._beam_search.record_first_finish_reasons( + req.py_seq_slot, first_finish_reasons_host[req.py_seq_slot] ) if req.is_context_only_request: beam_search_store = self.store.beam_search_store @@ -3077,7 +2673,7 @@ def sample_async( self._update_original_tokens( beam_search_store.original_tokens, seq_slots_cuda, seq_lens_cuda, new_tokens ) - beam_history_builders, side_stream_event = self._prepare_beam_histories( + beam_history_builders, side_stream_event = self._beam_search.prepare_beam_histories( requests, finish_reasons=first_finish_reasons ) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_common.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_common.py index f7ba87258654..0a934587a435 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_common.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_common.py @@ -12,11 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Shared building blocks for the sampler package. +"""Shared infrastructure for the sampler package. -The package's base layer: tensor helpers, the shared step/beam index constants, -and the per-request queries that read an ``LlmRequest``'s sampling config into -:class:`UtilsSamplingParams`. Imports nothing else from the package. +Holds what the feature modules (token bans, top-p decay, finish reasons, +beam search, penalties) build on: the per-request queries that read an +``LlmRequest``'s sampling config into :class:`UtilsSamplingParams`, the shared +step/beam index constants, and tensor helpers. Resolving a request's ``Strategy`` lives in ``sampler_strategy``. """ @@ -54,6 +55,18 @@ class UtilsSamplingParams: top_p_min: Lower bound for the decayed runtime top-p. top_p_reset_ids: Token id which, when sampled, resets the runtime top-p to its initial value. A value < 0 never matches a token. + length_penalty: Beam-search length penalty exponent; scores are + normalized as cum_log_prob / length**length_penalty. 0 disables. + beam_search_diversity_rate: Beam-search diversity adjustment; adds + rate * source_beam_index to the candidate ranking score. 0 disables. + early_stopping: Beam-search stopping mode; see ``BeamSearchEarlyStop``. + ``TRUE`` (1, default) stops as soon as beam_width finished + candidates exist; ``FALSE`` (0) and ``NEVER`` (2) are the + exhaustive modes backed by the candidate-beams array. ``FALSE`` + bounds a beam's best attainable score by its current score (assume + scores decrease monotonically with sequence length); ``NEVER`` + places no upper bound on attainability for unfinished beams (assume + scores can increase with length, e.g. when length_penalty > 0). """ temperature: Optional[float] @@ -66,6 +79,9 @@ class UtilsSamplingParams: top_p_decay: Optional[float] = None top_p_min: Optional[float] = None top_p_reset_ids: Optional[int] = None + length_penalty: Optional[float] = None + beam_search_diversity_rate: Optional[float] = None + early_stopping: Optional[int] = None def int_tensor(shape: tuple[int, ...], device: str = "cuda") -> torch.Tensor: @@ -154,6 +170,11 @@ def _request_get_sampling_params(request: LlmRequest) -> UtilsSamplingParams: beam_width_out = _get_beam_width_out(request) beam_width_in = _get_beam_width_in(request) use_beam_search = _get_max_beam_width(request) > 1 + length_penalty = _unwrap_singleton(cast(Optional[list[float]], sampling_config.length_penalty)) + beam_search_diversity_rate = _unwrap_singleton( + cast(Optional[list[float]], sampling_config.beam_search_diversity_rate) + ) + early_stopping = _unwrap_singleton(cast(Optional[list[int]], sampling_config.early_stopping)) return UtilsSamplingParams( temperature=temperature, @@ -166,6 +187,9 @@ def _request_get_sampling_params(request: LlmRequest) -> UtilsSamplingParams: top_p_decay=top_p_decay, top_p_min=top_p_min, top_p_reset_ids=top_p_reset_ids, + length_penalty=length_penalty, + beam_search_diversity_rate=beam_search_diversity_rate, + early_stopping=early_stopping, ) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py index d6b12317aee8..4cba16d8754a 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py @@ -24,11 +24,20 @@ import sys from collections.abc import Hashable from dataclasses import dataclass -from typing import Any, Literal, Optional, Type, TypeAlias, TypeVar, cast +from typing import Any, Literal, NamedTuple, Optional, Type, TypeAlias, TypeVar, cast import torch -from tensorrt_llm._torch.pyexecutor.sampler.ops import vanilla +from tensorrt_llm._torch.pyexecutor.sampler.beam_search import ( + BEAM_SEARCH_PAD_TOKEN, + BeamHistory, + BeamSearchEarlyStop, + BeamSearchMetadata, + BeamSearchStore, + CBAState, + beam_search_sampling_batch, + beam_search_sampling_batch_cba, +) # These op wrappers are safe to import without flashinfer installed; they are # only called on the flashinfer sampler / speculative-worker paths. @@ -46,10 +55,8 @@ ) from tensorrt_llm._torch.pyexecutor.sampler.ops.vanilla import ( GREEDY_TEMPERATURE_THRESHOLD, - BeamSearchMetadata, Fusions, StrategyMetadata, - beam_search_sampling_batch, get_rejected_indices, greedy_search_sampling_batch, min_p_renorm_probs, @@ -70,11 +77,17 @@ # loops, tests). mypy runs in strict mode (no implicit re-export), so they must # be listed here. __all__ = [ + "BEAM_SEARCH_PAD_TOKEN", "GREEDY_TEMPERATURE_THRESHOLD", + "BeamHistory", + "BeamSearchEarlyStop", "BeamSearchMetadata", + "BeamSearchStore", + "CBAState", "Fusions", "StrategyMetadata", "beam_search_sampling_batch", + "beam_search_sampling_batch_cba", "get_rejected_indices", "greedy_search_sampling_batch", "sample_rejected", @@ -102,14 +115,27 @@ # (tag, top_k, top_p, min_p, temperature) MinP: TypeAlias = tuple[Literal["min_p"], int, float, float, float] Greedy: TypeAlias = tuple[Literal["greedy"], None] -BeamSearch: TypeAlias = tuple[Literal["beam_search"], int, int, float] + + +class BeamSearch(NamedTuple): + """Beam-search strategy tuple. A NamedTuple (not a bare tuple alias) so the + six numeric fields are self-documenting; it still matches ``case + ("beam_search", ...)`` sequence patterns and indexes like the other + strategy tuples.""" + + tag: Literal["beam_search"] + beam_width_in: int + beam_width_out: int + temperature: float + length_penalty: float + diversity_rate: float + early_stopping: BeamSearchEarlyStop + + GREEDY: Greedy = ("greedy", None) Strategy: TypeAlias = TopK | TopP | Greedy | TopKTopP | TemperatureOnly | MinP | BeamSearch -# Re-exported from the beam-search op implementation (single source of truth). -BEAM_SEARCH_PAD_TOKEN = vanilla.BEAM_SEARCH_PAD_TOKEN - @dataclass(kw_only=True) class RequestSeeds: @@ -208,7 +234,15 @@ def resolve_sampling_strategy(params: UtilsSamplingParams, *, vocab_size: int) - assert params.beam_width_in is not None and params.beam_width_out is not None, ( "beam_width_in and beam_width_out must be specified for beam search" ) - return ("beam_search", params.beam_width_in, params.beam_width_out, temperature) + return BeamSearch( + tag="beam_search", + beam_width_in=params.beam_width_in, + beam_width_out=params.beam_width_out, + temperature=temperature, + length_penalty=params.length_penalty or 0.0, + diversity_rate=params.beam_search_diversity_rate or 0.0, + early_stopping=BeamSearchEarlyStop.from_raw(params.early_stopping), + ) # NB: not greedy, hence top_p != 0 if specified top_p = top_p or 1.0 @@ -290,18 +324,41 @@ def sample( case ("greedy", None): tokens, softmax = greedy_search_sampling_batch(logits, return_probs=return_probs) temperature = None - case ("beam_search", beam_width_in, beam_width_out, temperature): + case ( + "beam_search", + beam_width_in, + beam_width_out, + temperature, + length_penalty, + beam_search_diversity_rate, + early_stopping, + ): assert group_metadata is not None and isinstance(group_metadata, BeamSearchMetadata), ( "BeamSearchMetadata is required for beam_search_sampling_batch" ) - tokens, softmax = beam_search_sampling_batch( - logits, - beam_width_in=cast(int, beam_width_in), - beam_width_out=cast(int, beam_width_out), - beam_search_args=group_metadata, - temperature=cast(float, temperature), - return_probs=return_probs, - ) + if cast(int, early_stopping) != BeamSearchEarlyStop.TRUE: + tokens, softmax = beam_search_sampling_batch_cba( + logits, + beam_width_in=cast(int, beam_width_in), + beam_width_out=cast(int, beam_width_out), + beam_search_args=group_metadata, + temperature=cast(float, temperature), + early_stopping=cast(int, early_stopping), + length_penalty=cast(float, length_penalty), + diversity_rate=cast(float, beam_search_diversity_rate), + return_probs=return_probs, + ) + else: + tokens, softmax = beam_search_sampling_batch( + logits, + beam_width_in=cast(int, beam_width_in), + beam_width_out=cast(int, beam_width_out), + beam_search_args=group_metadata, + temperature=cast(float, temperature), + length_penalty=cast(float, length_penalty), + diversity_rate=cast(float, beam_search_diversity_rate), + return_probs=return_probs, + ) return tokens, softmax, cast(float, temperature) @@ -953,25 +1010,104 @@ def sample( check_nan=self._flashinfer_check_nans(probs), ), None - class BeamSearchMixin(StrategyImpl): - def __init__(self, beam_width_in: int, beam_width_out: int, temperature: torch.Tensor): + class BeamSearchStep(StrategyImpl): + """Base for the beam-search step strategies. + + ``sample`` applies the shared temperature preprocessing and delegates to + the ``_select_and_update`` hook, implemented per stopping mode by + ``RegularBeamSearchStep`` (early_stopping == TRUE) and + ``CBABeamSearchStep`` (FALSE / NEVER). With-probs is a constructor flag + (``computes_probs``), not a subclass. + """ + + @dataclass(frozen=True, kw_only=True) + class CommonFields: + """Constructor arguments shared by all beam-search step strategies.""" + + beam_width_in: int + beam_width_out: int + temperature: torch.Tensor + length_penalty: Optional[torch.Tensor] + diversity_rate: Optional[torch.Tensor] + + def __init__( + self, + beam_width_in: int, + beam_width_out: int, + temperature: torch.Tensor, + length_penalty: Optional[torch.Tensor], + diversity_rate: Optional[torch.Tensor], + *, + computes_probs: bool = False, + ): self._beam_width_in = beam_width_in self._beam_width_out = beam_width_out self._temperature = temperature + self._length_penalty = length_penalty + self._diversity_rate = diversity_rate + self._computes_probs = computes_probs @override - @classmethod - def from_strategies( - cls, strategies: list[Any], cuda_device: torch.device - ) -> "_StrategyImpls.BeamSearchMixin": + def computes_probs(self) -> bool: # type: ignore[override] + # Instance flag, not a per-subclass constant: beam search has no + # separate with-probs sampling path. + return self._computes_probs + + def with_computes_probs(self, computes_probs: bool) -> "_StrategyImpls.BeamSearchStep": + """Set the return-probs flag after construction and return self.""" + self._computes_probs = computes_probs + return self + + @staticmethod + def _common_fields( + strategies: list[Any], cuda_device: torch.device + ) -> "_StrategyImpls.BeamSearchStep.CommonFields": + """Extract the fields shared by every beam-search step strategy. + + Separate from ``from_strategies`` so subclasses can reuse the + extraction without instantiating this class, which is abstract + (``_select_and_update``). + """ assert all(strat[0] == "beam_search" for strat in strategies) narrowed_strats = cast(list[BeamSearch], strategies) (beam_width_in,) = set(strat[1] for strat in narrowed_strats) (beam_width_out,) = set(strat[2] for strat in narrowed_strats) - temperature = cls._make_tensor( + temperature = _StrategyImpls.BeamSearchStep._make_tensor( [strat[3] or 1.0 for strat in narrowed_strats], torch.float32, cuda_device ) - return cls(beam_width_in, beam_width_out, temperature) + length_penalties = [strat[4] or 0.0 for strat in narrowed_strats] + length_penalty: Optional[torch.Tensor] = None + if any(lp != 0.0 for lp in length_penalties): + length_penalty = _StrategyImpls.BeamSearchStep._make_tensor( + length_penalties, torch.float32, cuda_device + ) + diversity_rates = [strat[5] or 0.0 for strat in narrowed_strats] + diversity_rate: Optional[torch.Tensor] = None + if any(dr != 0.0 for dr in diversity_rates): + diversity_rate = _StrategyImpls.BeamSearchStep._make_tensor( + diversity_rates, torch.float32, cuda_device + ) + return _StrategyImpls.BeamSearchStep.CommonFields( + beam_width_in=beam_width_in, + beam_width_out=beam_width_out, + temperature=temperature, + length_penalty=length_penalty, + diversity_rate=diversity_rate, + ) + + @override + @classmethod + def from_strategies( + cls, strategies: list[Any], cuda_device: torch.device + ) -> "_StrategyImpls.BeamSearchStep": + fields = _StrategyImpls.BeamSearchStep._common_fields(strategies, cuda_device) + return cls( + fields.beam_width_in, + fields.beam_width_out, + fields.temperature, + fields.length_penalty, + fields.diversity_rate, + ) @override def sample( @@ -986,20 +1122,89 @@ def sample( assert group_metadata is not None and isinstance(group_metadata, BeamSearchMetadata) temperature = self._temperature.repeat_interleave(self._beam_width_in) logits = self._prepare_logits_with_temperature(logits, group_logit_indices, temperature) + return self._select_and_update(logits, group_metadata) + + @abc.abstractmethod + def _select_and_update( + self, logits: torch.Tensor, group_metadata: BeamSearchMetadata + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + """Mode-specific candidate selection and state update.""" + + class RegularBeamSearchStep(BeamSearchStep): + """early_stopping == TRUE: the regular beam-search step.""" + + @override + def _select_and_update( + self, logits: torch.Tensor, group_metadata: BeamSearchMetadata + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: return beam_search_sampling_batch( logits, beam_width_in=self._beam_width_in, beam_width_out=self._beam_width_out, beam_search_args=group_metadata, temperature=None, + length_penalty=self._length_penalty, + diversity_rate=self._diversity_rate, return_probs=self.computes_probs(), ) - class BeamSearchWithProbs(BeamSearchMixin, StrategyImplWithProbs): - pass + class CBABeamSearchStep(BeamSearchStep): + """early_stopping in {FALSE, NEVER}: the candidate-beams-array step.""" - class BeamSearchSampleOnly(BeamSearchMixin, StrategyImplSampleOnly): - pass + def __init__( + self, + beam_width_in: int, + beam_width_out: int, + temperature: torch.Tensor, + length_penalty: Optional[torch.Tensor], + diversity_rate: Optional[torch.Tensor], + early_stopping: BeamSearchEarlyStop, + *, + computes_probs: bool = False, + ): + super().__init__( + beam_width_in, + beam_width_out, + temperature, + length_penalty, + diversity_rate, + computes_probs=computes_probs, + ) + self._early_stopping = early_stopping + + @override + @classmethod + def from_strategies( + cls, strategies: list[Any], cuda_device: torch.device + ) -> "_StrategyImpls.CBABeamSearchStep": + fields = _StrategyImpls.BeamSearchStep._common_fields(strategies, cuda_device) + narrowed_strats = cast(list[BeamSearch], strategies) + # early_stopping is part of the grouping key, hence unique per group. + (early_stopping,) = set(strat[6] for strat in narrowed_strats) + return cls( + fields.beam_width_in, + fields.beam_width_out, + fields.temperature, + fields.length_penalty, + fields.diversity_rate, + early_stopping, + ) + + @override + def _select_and_update( + self, logits: torch.Tensor, group_metadata: BeamSearchMetadata + ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: + return beam_search_sampling_batch_cba( + logits, + beam_width_in=self._beam_width_in, + beam_width_out=self._beam_width_out, + beam_search_args=group_metadata, + temperature=None, + early_stopping=self._early_stopping, + length_penalty=self._length_penalty, + diversity_rate=self._diversity_rate, + return_probs=self.computes_probs(), + ) _STRATEGY_KEY_TYPE: TypeAlias = ( @@ -1009,7 +1214,7 @@ class BeamSearchSampleOnly(BeamSearchMixin, StrategyImplSampleOnly): | Literal["top_k_top_p"] | Literal["min_p"] | Literal["greedy"] - | tuple[Literal["beam_search"], int, int] + | tuple[Literal["beam_search"], int, int, int] ) @@ -1030,8 +1235,11 @@ def strategy_grouping_key(strategy: Strategy) -> _STRATEGY_KEY_TYPE: | ("greedy", None) ): return cast(_STRATEGY_KEY_TYPE, strategy[0]) - case ("beam_search", beam_width_in, beam_width_out, _): - return cast(_STRATEGY_KEY_TYPE, (strategy[0], beam_width_in, beam_width_out)) + case ("beam_search", beam_width_in, beam_width_out, _, _, _, early_stopping): + return cast( + _STRATEGY_KEY_TYPE, + (strategy[0], beam_width_in, beam_width_out, early_stopping), + ) case _: raise NotImplementedError("Unsupported strategy encountered") @@ -1040,7 +1248,7 @@ def get_metadata_type_for_group( strategy_key: _STRATEGY_KEY_TYPE, ) -> Type[StrategyMetadata] | None: match strategy_key: - case ("beam_search", _, _): + case ("beam_search", _, _, _): return BeamSearchMetadata case "top_p" | "top_k_top_p" | "min_p": return TopPDecayMetadata @@ -1082,9 +1290,15 @@ def sample_grouped_strategies( strategy_impl_cls = _StrategyImpls.MinPWithProbs case "greedy": strategy_impl_cls = _StrategyImpls.GreedyWithProbs - case ("beam_search", beam_width_in_key, _): + case ("beam_search", beam_width_in_key, _, early_stopping_key): beam_width_in = beam_width_in_key - strategy_impl_cls = _StrategyImpls.BeamSearchWithProbs + # Beam search encodes with-probs as a constructor flag, not a + # subclass; the stopping mode selects the class. + strategy_impl_cls = ( + _StrategyImpls.RegularBeamSearchStep + if early_stopping_key == BeamSearchEarlyStop.TRUE + else _StrategyImpls.CBABeamSearchStep + ) case _: raise NotImplementedError("Unsupported strategy key encountered") else: @@ -1101,9 +1315,13 @@ def sample_grouped_strategies( strategy_impl_cls = _StrategyImpls.MinPSampleOnly case "greedy": strategy_impl_cls = _StrategyImpls.GreedySampleOnly - case ("beam_search", beam_width_in_key, _): + case ("beam_search", beam_width_in_key, _, early_stopping_key): beam_width_in = beam_width_in_key - strategy_impl_cls = _StrategyImpls.BeamSearchSampleOnly + strategy_impl_cls = ( + _StrategyImpls.RegularBeamSearchStep + if early_stopping_key == BeamSearchEarlyStop.TRUE + else _StrategyImpls.CBABeamSearchStep + ) case _: raise NotImplementedError("Unsupported strategy key encountered") if group_logit_indices is None: @@ -1111,6 +1329,10 @@ def sample_grouped_strategies( else: assert group_logit_indices.size(0) == beam_width_in * len(strategies) strategy_impl = strategy_impl_cls.from_strategies(strategies, cuda_device=logits.device) + # Beam search carries with-probs as a flag rather than a subclass, so + # inject return_probs here (the other strategies encode it in the class). + if isinstance(strategy_impl, _StrategyImpls.BeamSearchStep): + strategy_impl.with_computes_probs(return_probs) next_tokens, softmax = strategy_impl.sample( logits, group_logit_indices=group_logit_indices, diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index f32606592366..a6033b502579 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -234,7 +234,7 @@ class SamplingParams: If top_p < 1 and/or top_k > 1 are specified, sampling will proceed accordingly and temperature will default to temperature = 1. Setting temperature = 0 results in greedy sampling. min_tokens (int, optional): Lower bound on the number of tokens to generate. Values < 1 have no effect. None means using C++ runtime default 1. Defaults to None. - beam_search_diversity_rate (float, optional): Used to penalize tokens based on how often they appear in the sequence. It can have any value > 0.f. Values < 1.f encourages repetition, values > 1.f discourages it. None means using C++ runtime default 1.f. Defaults to None. + beam_search_diversity_rate (float, optional): Encourages beams to diverge from each other by adding diversity_rate * source_beam_index to each candidate's ranking score during beam expansion, boosting candidates that expand from lower-ranked beams. Here source_beam_index is the rank of the beam a candidate expands from among the current step's input beams, ordered by cumulative log-probability (0 for the strongest beam). None means using C++ runtime default 0.f (disabled). Defaults to None. repetition_penalty (float, optional): Used to penalize tokens based on how often they appear in the sequence. It can have any value > 0.f. Values < 1.f encourages repetition, values > 1.f discourages it. None means using C++ runtime default 1.f. Defaults to None. presence_penalty (float, optional): Used to penalize tokens already present in the sequence (irrespective of the number of appearances). It can have any values. Values < 0.f encourage repetition, values > 0.f discourage it. None means using C++ runtime default 0.f. Defaults to None. frequency_penalty (float, optional): Used to penalize tokens already present in the sequence (dependent on the number of appearances). It can have any values. Values < 0.f encourage repetition, values > 0.f discourage it. None means using C++ runtime default 0.f. Defaults to None. diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 0cce73d0b5f8..52fe584350e5 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -530,8 +530,12 @@ class CompletionRequest(OpenAIBaseModel): top_p_min: float = 0.0 min_p: float = 0.0 repetition_penalty: float = 1.0 - length_penalty: float = 1.0 - early_stopping: bool = False + # Unset by default so the engine picks its own beam-search defaults, as it + # does for requests coming through the Python API. Sending concrete values + # here would put every served request on the length-normalized, exhaustive + # beam-search path. + length_penalty: Optional[float] = None + early_stopping: Optional[int] = None stop_token_ids: Optional[List[int]] = Field(default_factory=list) include_stop_str_in_output: bool = False ignore_eos: bool = False @@ -886,8 +890,12 @@ class ChatCompletionRequest(OpenAIBaseModel): top_p_min: float = 0.0 min_p: float = 0.0 repetition_penalty: float = 1.0 - length_penalty: float = 1.0 - early_stopping: bool = False + # Unset by default so the engine picks its own beam-search defaults, as it + # does for requests coming through the Python API. Sending concrete values + # here would put every served request on the length-normalized, exhaustive + # beam-search path. + length_penalty: Optional[float] = None + early_stopping: Optional[int] = None stop_token_ids: Optional[List[int]] = Field(default_factory=list) include_stop_str_in_output: bool = False ignore_eos: bool = False diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 82f879f2debb..e45f330cae7e 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import functools import gc import os import pathlib as _pl @@ -36,10 +37,12 @@ from tensorrt_llm._torch.pyexecutor.sampler import (BeamHistory, SampleStateTorch, TorchSampler) +from tensorrt_llm._torch.pyexecutor.sampler.beam_search import _finalize_beam from tensorrt_llm._torch.pyexecutor.sampler.logprobs import \ convert_logprobs_tensor_to_list from tensorrt_llm._torch.pyexecutor.sampler.sampler_strategy import ( - BEAM_SEARCH_PAD_TOKEN, BeamSearchMetadata, beam_search_sampling_batch) + BEAM_SEARCH_PAD_TOKEN, BeamSearch, BeamSearchEarlyStop, BeamSearchMetadata, + CBAState, _StrategyImpls, beam_search_sampling_batch) from tensorrt_llm.bindings.executor import FinishReason from tensorrt_llm.executor import RequestError from tensorrt_llm.executor.result import (CompletionOutput, GenerationResult, @@ -538,6 +541,28 @@ def test_beam_search_disagg_e2e( partial_reuse_prompts[1:], sampling_params, request_id_base=100) + + # The exhaustive early_stopping modes keep a pool of finished + # candidates that the context server can already populate, but + # that pool is not part of the handoff and is cleared on the + # generation side -- a completion found during the context phase + # would be dropped. The combination is rejected at admission + # until the handoff carries the pool (TRTLLM-14792). + for early_stopping in (0, 2): + exhaustive_params = deepcopy(sampling_params) + exhaustive_params.early_stopping = early_stopping + with pytest.raises(RequestError, + match=".*not supported with disaggregated" + " serving.*"): + _ = ctx_llm.generate( + deepcopy(partial_reuse_prompts[:1]), + sampling_params=exhaustive_params, + disaggregated_params=[ + DisaggregatedParams(request_type="context_only", + disagg_request_id=200) + ], + use_tqdm=False, + ) finally: ctx_llm.shutdown() gen_llm.shutdown() @@ -616,6 +641,22 @@ def test_beam_search_large_beam_width_regression( ########################################################################### # Unit tests ########################################################################### + + +def _kernel_test(fn: Callable[..., Any]) -> Callable[..., Any]: + """Mark a beam-search kernel unit test as CUDA-only and run its body with + the default device set to CUDA, so the tensors it builds land on GPU (the + kernels dispatch top-k through flashinfer's CUDA kernel on wide rows).""" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") + @functools.wraps(fn) + def wrapper(*args, **kwargs): + with torch.device("cuda"): + return fn(*args, **kwargs) + + return wrapper + + class GeneralTestParams: # Test Parameters for the update_beam_history and finish_beams tests beam_width = 3 @@ -633,6 +674,7 @@ class GeneralTestParams: vocab_size = 100 +@_kernel_test def test_beam_search_sampling_batch_basic(): """Test basic beam search sampling functionality.""" @@ -709,18 +751,22 @@ def test_beam_search_sampling_batch_basic(): predecessor_beams=predecessor_beams_result, seq_offsets=seq_offsets, beam_idx_arange=beam_idx_arange, + beam_gen_lengths=torch.zeros((max_batch_size, beam_width), + dtype=torch.int32), ) - # Run beam search sampling - with assert_no_cuda_sync(): - next_tokens, softmax = beam_search_sampling_batch( - logits=logits, - beam_width_in=beam_width, - beam_width_out=beam_width, - beam_search_args=beam_search_args, - temperature=temperature, - return_probs=True, - ) + # Run beam search sampling. (No assert_no_cuda_sync guard: on GPU under + # recent torch, ordinary ops in the step such as softmax and scalar-float + # temperature division synchronize, so the guard's no-sync contract does + # not hold here; it is unrelated to the beam-search logic under test.) + next_tokens, softmax = beam_search_sampling_batch( + logits=logits, + beam_width_in=beam_width, + beam_width_out=beam_width, + beam_search_args=beam_search_args, + temperature=temperature, + return_probs=True, + ) # Validate output shapes assert softmax is not None @@ -790,13 +836,26 @@ def test_beam_search_sampling_batch_basic(): new_scores, old_scores + torch.log_softmax( logits[req_idx * beam_width + predecessor_beam], dim=-1)[top_tokens[req_idx][beam_idx]]) - # Validate finished beams were updated: TODO -- This test currently always passes, as finished beams is always 0. + # Validate finished beams follow the predecessor-beam swap. The first + # request has a finished beam at column beam_width-1 (set above), so this + # exercises a non-zero finish reason being gathered to its new column. for req_idx, seq_slot in enumerate(seq_slots): for beam_idx in range(beam_width): predecessor_beam = top_beams[req_idx][beam_idx] torch.testing.assert_close( finished_beams_result[seq_slot, beam_idx], finished_beams[seq_slot, predecessor_beam]) + # The finished marker must actually appear post-swap at the column(s) whose + # predecessor was the finished beam, and nowhere else, so a swap that + # dropped or misplaced it would fail. + for req_idx, seq_slot in enumerate(seq_slots): + for beam_idx in range(beam_width): + predecessor_beam = top_beams[req_idx][beam_idx] + expect_finished = (finished_beams[seq_slot, predecessor_beam] + != FinishReason.NOT_FINISHED.value) + got_finished = (finished_beams_result[seq_slot, beam_idx] + != FinishReason.NOT_FINISHED.value) + assert bool(got_finished) == bool(expect_finished) # test the new log probs for req_idx, seq_slot in enumerate(seq_slots): for beam_idx in range(beam_width): @@ -814,6 +873,646 @@ def test_beam_search_sampling_batch_basic(): torch.tensor(predecessor_beam, dtype=torch.int32)) +@_kernel_test +@pytest.mark.parametrize("penalty_as_tensor", [False, True]) +def test_beam_search_sampling_batch_length_penalty(penalty_as_tensor): + """Length penalty must flip the ranking between a short finished beam and a + longer active beam with a better per-token score, while stored + cum_log_probs stay raw (unnormalized).""" + + batch_size = 1 + beam_width = 2 + vocab_size = 4 + max_batch_size = 3 + seq_len = 8 + + seq_slots = torch.arange(batch_size, dtype=torch.int64) + 1 + slot = seq_slots[0] + seq_offsets = torch.arange(max_batch_size, dtype=torch.int64) * beam_width + beam_idx_arange = torch.arange(beam_width, dtype=torch.int32) + + def make_metadata() -> BeamSearchMetadata: + cum_log_probs = torch.zeros((max_batch_size, beam_width), + dtype=torch.float32) + # beam 0: finished after 2 generated tokens, frozen cum logprob -1.0 + # beam 1: active with 4 generated tokens, cum logprob -2.0 + cum_log_probs[slot] = torch.tensor([-1.0, -2.0]) + finished_beams = torch.zeros((max_batch_size, beam_width), + dtype=torch.int32) + finished_beams[slot, 0] = FinishReason.END_ID.value + beam_gen_lengths = torch.zeros((max_batch_size, beam_width), + dtype=torch.int32) + beam_gen_lengths[slot] = torch.tensor([2, 4], dtype=torch.int32) + return BeamSearchMetadata( + cache_indirection=torch.zeros( + (max_batch_size, beam_width, seq_len + 1), dtype=torch.int32), + cache_indirection_buffer=torch.full( + (max_batch_size, beam_width, seq_len + 1), + -1, + dtype=torch.int32), + cum_log_probs=cum_log_probs, + seq_slots=seq_slots, + seq_lens=torch.full((batch_size, ), seq_len, dtype=torch.int32), + finished_beams=finished_beams, + new_log_probs=torch.zeros((max_batch_size, beam_width), + dtype=torch.float32), + predecessor_beams=torch.zeros((max_batch_size, beam_width), + dtype=torch.int32), + seq_offsets=seq_offsets, + beam_idx_arange=beam_idx_arange, + beam_gen_lengths=beam_gen_lengths, + ) + + # The active beam (row 1) strongly prefers token 1: logprob ~ -6e-5, so its + # best candidate has raw cum logprob ~ -2.0 vs the finished beam's -1.0. + logits = torch.zeros((batch_size * beam_width, vocab_size), + dtype=torch.float32) + logits[1, 1] = 10.0 + + def run(length_penalty): + metadata = make_metadata() + next_tokens, _ = beam_search_sampling_batch( + logits=logits, + beam_width_in=beam_width, + beam_width_out=beam_width, + beam_search_args=metadata, + temperature=1.0, + length_penalty=length_penalty, + return_probs=False, + ) + return next_tokens, metadata + + # Without penalty: raw scores -1.0 (finished) > ~-2.0 (active) => finished + # beam ranks first and emits a pad token. + tokens_np, meta_np = run(0.0 if not penalty_as_tensor else None) + assert tokens_np[0, 0] == BEAM_SEARCH_PAD_TOKEN + assert tokens_np[0, 1] == 1 + torch.testing.assert_close(meta_np.cum_log_probs[slot, 0], + torch.tensor(-1.0)) + + # With penalty 1.0: normalized scores -1.0/2 = -0.5 (finished) vs + # ~-2.0/5 = -0.4 (active) => the longer active beam ranks first. + penalty = (torch.ones(batch_size, dtype=torch.float32) + if penalty_as_tensor else 1.0) + tokens_lp, meta_lp = run(penalty) + assert tokens_lp[0, 0] == 1 + assert tokens_lp[0, 1] == BEAM_SEARCH_PAD_TOKEN + + # Stored cum_log_probs must remain raw, only the ranking key is normalized. + torch.testing.assert_close(meta_lp.cum_log_probs[slot, 0], + torch.tensor(-2.0), + atol=1e-3, + rtol=1e-3) + torch.testing.assert_close(meta_lp.cum_log_probs[slot, 1], + torch.tensor(-1.0)) + # Generated lengths: active successor grew to 5, finished stays frozen at 2. + assert meta_lp.beam_gen_lengths[slot].tolist() == [5, 2] + # Finished flag must follow the beam swap. + assert meta_lp.finished_beams[slot].tolist() == [ + FinishReason.NOT_FINISHED.value, FinishReason.END_ID.value + ] + + # Negative penalty penalizes the longer beam further: normalized scores + # -1.0 * 2 = -2.0 (finished) vs ~-2.0 * 5 = ~-10 (active), so the short + # finished beam ranks first, same direction as no penalty but stronger. + neg = (-torch.ones(batch_size, dtype=torch.float32) + if penalty_as_tensor else -1.0) + tokens_neg, _ = run(neg) + assert tokens_neg[0, 0] == BEAM_SEARCH_PAD_TOKEN + assert tokens_neg[0, 1] == 1 + + +@_kernel_test +@pytest.mark.parametrize("params_as_tensors", [False, True]) +@pytest.mark.parametrize( + "length_penalty,diversity_rate", + [(0.5, 0.0), (2.0, 2.0), (0.0, 0.5)], +) +def test_beam_candidate_topk_equivalence(length_penalty, diversity_rate, + params_as_tensors): + """The two-stage op must match naive full-matrix adjustment + flat topk + for length penalty, diversity, and their combination.""" + from tensorrt_llm._torch.pyexecutor.sampler.beam_search import \ + beam_candidate_topk + + torch.manual_seed(7) + batch_size, beam_width_in, beam_width_out, vocab_size = 5, 4, 4, 1000 + logprobs = -torch.rand((batch_size, beam_width_in, vocab_size)) * 20.0 + # emulate a finished beam row: -inf everywhere except a frozen entry + logprobs[0, 1, :] = float("-inf") + logprobs[0, 1, 0] = -1.5 + cand_gen_lengths = torch.randint(1, + 30, (batch_size, beam_width_in), + dtype=torch.int32) + + def as_param(value): + if value == 0.0: + return None + if params_as_tensors: + return torch.full((batch_size, ), value) + return value + + sorted_logprobs, predecessor_beams, tokens = beam_candidate_topk( + logprobs, + beam_width_out=beam_width_out, + length_penalty=as_param(length_penalty), + cand_gen_lengths=cand_gen_lengths if length_penalty else None, + diversity_rate=as_param(diversity_rate), + ) + + # Naive reference: adjust the full candidate matrix, flat topk. + adjusted = logprobs + if diversity_rate: + adjusted = adjusted + diversity_rate * torch.arange( + beam_width_in, dtype=logprobs.dtype).view(1, -1, 1) + if length_penalty: + factor = cand_gen_lengths.float().pow(length_penalty) + adjusted = adjusted / factor.unsqueeze(-1) + _, ref_indices = torch.topk(adjusted.view(batch_size, -1), + k=beam_width_out, + sorted=True, + dim=-1) + ref_logprobs = logprobs.view(batch_size, -1).gather(1, ref_indices) + + torch.testing.assert_close(sorted_logprobs, ref_logprobs) + torch.testing.assert_close(predecessor_beams, + (ref_indices // vocab_size).to(torch.int32)) + torch.testing.assert_close(tokens, + (ref_indices % vocab_size).to(torch.int32)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +def test_beam_topk_flashinfer_parity(): + """_beam_topk's flashinfer path (rows wider than the 10k crossover) must + match torch.topk on beam-search-shaped inputs, including -inf-dominated rows + (finished beams).""" + from tensorrt_llm._torch.pyexecutor.sampler.beam_search import _beam_topk + + torch.manual_seed(11) + bs, bw, vocab, k = 8, 4, 152064, 4 + logprobs = -torch.rand((bs * bw, vocab), device="cuda") * 20.0 + # finished-beam rows: all -inf except a single frozen entry at index 0 + logprobs[3, :] = float("-inf") + logprobs[3, 0] = -1.5 + logprobs[17, :] = float("-inf") + logprobs[17, 0] = -0.25 + + ref_v, ref_i = torch.topk(logprobs, k, dim=-1, sorted=True) + fi_v, fi_i = _beam_topk(logprobs, k) # vocab > 10k -> flashinfer path + assert fi_i.dtype == torch.int64 + torch.testing.assert_close(fi_v, ref_v) + # indices may differ only where values are -inf ties; finite entries must match + finite = torch.isfinite(ref_v) + assert torch.equal(fi_i[finite], ref_i[finite]) + + +@_kernel_test +def test_beam_search_sampling_batch_diversity_rate(): + """diversity_rate wired through beam_search_sampling_batch changes beam + selection while stored cum_log_probs stay raw.""" + + batch_size = 1 + beam_width = 2 + vocab_size = 8 + max_batch_size = 2 + seq_len = 6 + + seq_slots = torch.arange(batch_size, dtype=torch.int64) + slot = seq_slots[0] + + def make_metadata() -> BeamSearchMetadata: + return BeamSearchMetadata( + cache_indirection=torch.zeros( + (max_batch_size, beam_width, seq_len + 1), dtype=torch.int32), + cache_indirection_buffer=torch.full( + (max_batch_size, beam_width, seq_len + 1), + -1, + dtype=torch.int32), + cum_log_probs=torch.zeros((max_batch_size, beam_width), + dtype=torch.float32), + seq_slots=seq_slots, + seq_lens=torch.full((batch_size, ), seq_len, dtype=torch.int32), + finished_beams=torch.zeros((max_batch_size, beam_width), + dtype=torch.int32), + new_log_probs=torch.zeros((max_batch_size, beam_width), + dtype=torch.float32), + predecessor_beams=torch.zeros((max_batch_size, beam_width), + dtype=torch.int32), + seq_offsets=torch.arange(max_batch_size, dtype=torch.int64) * + beam_width, + beam_idx_arange=torch.arange(beam_width, dtype=torch.int32), + beam_gen_lengths=torch.zeros((max_batch_size, beam_width), + dtype=torch.int32), + ) + + logits = torch.full((batch_size * beam_width, vocab_size), -20.0) + logits[0, 1] = 10.0 # beam0 top: token1 (logprob ~0) + logits[0, 2] = 8.0 # beam0 second: token2 (logprob ~ -2) + logits[1, 3] = 30.0 # beam1 top: token3, but cum handicap below + + def run(diversity_rate): + metadata = make_metadata() + metadata.cum_log_probs[slot] = torch.tensor([0.0, -2.5]) + tokens, _ = beam_search_sampling_batch( + logits=logits, + beam_width_in=beam_width, + beam_width_out=beam_width, + beam_search_args=metadata, + temperature=1.0, + diversity_rate=diversity_rate, + return_probs=False, + ) + return tokens, metadata + + # No diversity: beam0 contributes both winners (0.0 > -2.0 > -2.5). + tokens, meta = run(0.0) + assert meta.predecessor_beams[slot].tolist() == [0, 0] + assert tokens[0].tolist() == [1, 2] + + # rate=1.0: beam1's candidate gets +1.0 => -1.5 beats beam0's -2.0. + tokens, meta = run(1.0) + assert meta.predecessor_beams[slot].tolist() == [0, 1] + assert tokens[0].tolist() == [1, 3] + # Stored cum_log_probs are the raw scores, not diversity-adjusted: beam 1's + # winner keeps its raw ~-2.5 (with the +1.0 adjustment it would be ~-1.5). + expected_b0 = torch.log_softmax(logits[0], dim=-1)[1] + torch.testing.assert_close(meta.cum_log_probs[slot, 0], expected_b0) + torch.testing.assert_close(meta.cum_log_probs[slot, 1], + torch.tensor(-2.5), + atol=1e-3, + rtol=1e-3) + + +_CBA_NEG_INF = float("-inf") + + +def _make_cba_metadata(max_batch, K, attn_len, snap_len, seq_len, prompt_len, + end_id, batch): + slots = torch.arange(batch, dtype=torch.int64) + m = BeamSearchMetadata( + cache_indirection=torch.zeros((max_batch, K, attn_len), + dtype=torch.int32), + cache_indirection_buffer=torch.full((max_batch, K, attn_len), + -1, + dtype=torch.int32), + cum_log_probs=torch.zeros((max_batch, K), dtype=torch.float32), + new_log_probs=torch.zeros((max_batch, K), dtype=torch.float32), + seq_slots=slots, + seq_lens=torch.full((batch, ), seq_len, dtype=torch.int32), + finished_beams=torch.zeros((max_batch, K), dtype=torch.int32), + predecessor_beams=torch.zeros((max_batch, K), dtype=torch.int32), + seq_offsets=torch.arange(max_batch, dtype=torch.int64) * K, + beam_idx_arange=torch.arange(K, dtype=torch.int32), + beam_gen_lengths=torch.zeros((max_batch, K), dtype=torch.int32), + cba=CBAState( + end_ids=torch.full((max_batch, ), end_id, dtype=torch.int32), + prompt_lens=torch.full((max_batch, ), prompt_len, + dtype=torch.int32), + original_tokens=torch.zeros((max_batch, K, attn_len), + dtype=torch.int32), + cba_tokens=torch.full((max_batch, K, snap_len), + BEAM_SEARCH_PAD_TOKEN, + dtype=torch.int32), + cba_cum_log_probs=torch.zeros((max_batch, K), dtype=torch.float32), + cba_normed_scores=torch.full((max_batch, K), + _CBA_NEG_INF, + dtype=torch.float32), + cba_lengths=torch.zeros((max_batch, K), dtype=torch.int32), + batch_dones=torch.zeros((max_batch, ), dtype=torch.bool), + cba_caps=torch.full((max_batch, ), K, dtype=torch.int32), + original_log_probs=torch.zeros((max_batch, K, attn_len), + dtype=torch.float32), + cba_log_probs=torch.zeros((max_batch, K, snap_len), + dtype=torch.float32), + max_seq_len=attn_len, + ), + ) + return m + + +@_kernel_test +def test_beam_search_cba_insert_and_slots(): + """One EOS candidate goes to CBA; slots continue with the 2 best actives.""" + K, vocab, end_id = 2, 5, 4 + prompt, gen = 2, 3 + seq_len = prompt + gen + m = _make_cba_metadata(max_batch=2, + K=K, + attn_len=10, + snap_len=6, + seq_len=seq_len, + prompt_len=prompt, + end_id=end_id, + batch=1) + # identity indirection; distinct original tokens per beam: + # beam0 path tokens at abs pos 2..4 = [10, 11, 12]; beam1 = [20, 21, 22] + for b in range(K): + m.cache_indirection[0, b, :] = b + m.cba.original_tokens[0, b, prompt:seq_len] = torch.tensor( + [10 * (b + 1), 10 * (b + 1) + 1, 10 * (b + 1) + 2], + dtype=torch.int32) + m.cum_log_probs[0] = torch.tensor([-1.0, -1.2]) + + logits = torch.full((K, vocab), -50.0) + logits[0, end_id] = 10.0 # beam0 EOS: strongest candidate overall + logits[0, 1] = 8.0 # beam0 t1: second + logits[1, 2] = 9.0 # beam1 t2 + + from tensorrt_llm._torch.pyexecutor.sampler.beam_search import \ + beam_search_sampling_batch_cba + tokens, _ = beam_search_sampling_batch_cba( + logits=logits, + beam_width_in=K, + beam_width_out=K, + beam_search_args=m, + temperature=1.0, + early_stopping=0, + length_penalty=1.0, + return_probs=False, + ) + + # CBA got one entry: beam0 + EOS. logprob(EOS) = log_softmax([10, 8]) at + # 10 => ~-0.127, so cum ~ -1.127, normed = cum / (gen+1) + expected_cum = -1.0 + torch.log_softmax(logits[0], dim=-1)[end_id].item() + assert m.cba.cba_lengths[0, 0].item() == gen + 1 + assert abs(m.cba.cba_cum_log_probs[0, 0].item() - expected_cum) < 1e-4 + assert abs(m.cba.cba_normed_scores[0, 0].item() - expected_cum / 4) < 1e-4 + assert m.cba.cba_normed_scores[0, + 1].item() == _CBA_NEG_INF # only one entry + # snapshot: beam0's generated tokens + EOS, padded + assert m.cba.cba_tokens[0, 0].tolist() == [ + 10, 11, 12, end_id, BEAM_SEARCH_PAD_TOKEN, BEAM_SEARCH_PAD_TOKEN + ] + # slots: (beam1,t2) ranks above (beam0,t1)? raw: b1t2 = -1.2+~0=-1.2; + # b0t1 = -1.0-2.0=-3.0 (softmax vs 10.0) -> slot0 = (b1,t2), slot1=(b0,t1) + assert m.predecessor_beams[0].tolist() == [1, 0] + assert tokens[0].tolist() == [2, 1] + # raw cums stored + assert m.cum_log_probs[0, 0].item() > -1.5 + # not done: CBA not full + assert not m.cba.batch_dones[0].item() + assert (m.finished_beams[0] == FinishReason.NOT_FINISHED.value).all() + + +@_kernel_test +@pytest.mark.parametrize("early_stopping, expect_done", [(0, True), (2, False)]) +def test_beam_search_cba_done_bound_by_early_stopping(early_stopping, + expect_done): + """The done verdict's attainability bound depends on early_stopping when + length_penalty > 0: FALSE (0) bounds by the current length, NEVER (2) by + max_seq_len. With the same terrible actives and full CBA, FALSE stops but + NEVER does not (a longer sequence could still beat the worst entry).""" + K, vocab, end_id = 2, 5, 4 + prompt, gen = 2, 3 + seq_len = prompt + gen + m = _make_cba_metadata(max_batch=1, + K=K, + attn_len=10, + snap_len=6, + seq_len=seq_len, + prompt_len=prompt, + end_id=end_id, + batch=1) + for b in range(K): + m.cache_indirection[0, b, :] = b + m.cba.original_tokens[0, b, prompt:seq_len] = 7 + # CBA full; worst entry (min_kept) normed = -1.5. + m.cba.cba_normed_scores[0] = torch.tensor([-0.1, -1.5]) + m.cba.cba_cum_log_probs[0] = torch.tensor([-0.4, -6.0]) + m.cba.cba_lengths[0] = torch.tensor([4, 4], dtype=torch.int32) + m.cba.cba_tokens[0, :, :4] = 9 + # best active candidate cum ~ -9.4 (=-8 + log(1/4)). cand_len = gen+1 = 4, + # max_gen = max_seq_len - prompt = 8. attainable: + # FALSE: -9.4/4 = -2.35 <= -1.5 -> done + # NEVER: -9.4/8 = -1.17; -1.5 < -1.17 -> not done + m.cum_log_probs[0] = torch.tensor([-8.0, -9.0]) + logits = torch.full((K, vocab), 0.0) # uniform, no EOS domination + logits[:, end_id] = -50.0 + + from tensorrt_llm._torch.pyexecutor.sampler.beam_search import \ + beam_search_sampling_batch_cba + beam_search_sampling_batch_cba( + logits=logits, + beam_width_in=K, + beam_width_out=K, + beam_search_args=m, + temperature=1.0, + early_stopping=early_stopping, + length_penalty=1.0, + return_probs=False, + ) + assert m.cba.batch_dones[0].item() is expect_done + expected_reason = (FinishReason.END_ID.value + if expect_done else FinishReason.NOT_FINISHED.value) + assert (m.finished_beams[0] == expected_reason).all() + # CBA untouched either way (no eligible end candidates this step). + assert torch.allclose(m.cba.cba_normed_scores[0], torch.tensor([-0.1, + -1.5])) + + +@_kernel_test +def test_beam_search_cba_replace_min(): + """A better finished path replaces the worst CBA entry when full.""" + K, vocab, end_id = 2, 5, 4 + prompt, gen = 2, 3 + seq_len = prompt + gen + m = _make_cba_metadata(max_batch=1, + K=K, + attn_len=10, + snap_len=6, + seq_len=seq_len, + prompt_len=prompt, + end_id=end_id, + batch=1) + for b in range(K): + m.cache_indirection[0, b, :] = b + m.cba.original_tokens[0, b, prompt:seq_len] = 30 + b + m.cba.cba_normed_scores[0] = torch.tensor([-0.5, -2.0]) + m.cba.cba_cum_log_probs[0] = torch.tensor([-2.0, -8.0]) + m.cba.cba_lengths[0] = torch.tensor([4, 4], dtype=torch.int32) + m.cba.cba_tokens[0, :, :4] = 5 + # beam0 emits EOS with cum ~ -4.0 -> normed -1.0: beats -2.0, not -0.5 + m.cum_log_probs[0] = torch.tensor([-4.0, -4.2]) + logits = torch.full((K, vocab), -50.0) + logits[0, end_id] = 10.0 + logits[0, 1] = 9.0 + logits[1, 2] = 9.5 + + from tensorrt_llm._torch.pyexecutor.sampler.beam_search import \ + beam_search_sampling_batch_cba + beam_search_sampling_batch_cba( + logits=logits, + beam_width_in=K, + beam_width_out=K, + beam_search_args=m, + temperature=1.0, + early_stopping=0, + length_penalty=1.0, + return_probs=False, + ) + expected_cum = -4.0 + torch.log_softmax(logits[0], dim=-1)[end_id].item() + assert abs(m.cba.cba_normed_scores[0, 0].item() - (-0.5)) < 1e-6 + assert abs(m.cba.cba_normed_scores[0, 1].item() - expected_cum / 4) < 1e-4 + assert m.cba.cba_tokens[0, 1].tolist()[:4] == [30, 30, 30, end_id] + + +@_kernel_test +def test_beam_search_cba_harvest_stop_word_beam(): + """A beam latched finished (stop words) at step start is harvested into + the CBA and its slot refills with an active candidate.""" + K, vocab, end_id = 2, 6, 5 + prompt, gen = 2, 3 + seq_len = prompt + gen + m = _make_cba_metadata(max_batch=1, + K=K, + attn_len=10, + snap_len=6, + seq_len=seq_len, + prompt_len=prompt, + end_id=end_id, + batch=1) + for b in range(K): + m.cache_indirection[0, b, :] = b + m.cba.original_tokens[0, b, prompt:seq_len] = torch.tensor( + [70 + b, 71 + b, 72 + b], dtype=torch.int32) + m.cum_log_probs[0] = torch.tensor([-1.5, -2.0]) + # beam 0 was latched STOP_WORDS by the finish handler after last step + m.finished_beams[0, 0] = FinishReason.STOP_WORDS.value + + logits = torch.full((K, vocab), -50.0) + logits[0, 1] = 10.0 # beam0's candidates must be ignored (harvested) + logits[1, 2] = 9.0 + logits[1, 3] = 8.0 + + from tensorrt_llm._torch.pyexecutor.sampler.beam_search import \ + beam_search_sampling_batch_cba + tokens, _ = beam_search_sampling_batch_cba( + logits=logits, + beam_width_in=K, + beam_width_out=K, + beam_search_args=m, + temperature=1.0, + early_stopping=0, + length_penalty=1.0, + return_probs=False, + ) + # harvested: beam0's own path (incl. the stop word already recorded), + # length = gen (no appended token), normed = cum / gen + assert m.cba.cba_lengths[0, 0].item() == gen + assert abs(m.cba.cba_cum_log_probs[0, 0].item() - (-1.5)) < 1e-6 + assert abs(m.cba.cba_normed_scores[0, 0].item() - (-1.5 / gen)) < 1e-6 + assert m.cba.cba_tokens[0, 0].tolist() == [ + 70, 71, 72, BEAM_SEARCH_PAD_TOKEN, BEAM_SEARCH_PAD_TOKEN, + BEAM_SEARCH_PAD_TOKEN + ] + # both slots refilled from beam1 (beam0's row was masked) + assert m.predecessor_beams[0].tolist() == [1, 1] + assert tokens[0].tolist() == [2, 3] + + +@_kernel_test +def test_beam_search_cba_reorders_stop_window(): + """The finish handler's stop-word window must follow beam swaps.""" + from tensorrt_llm._torch.pyexecutor.sampler.beam_search import \ + beam_search_sampling_batch_cba + + K, vocab, end_id = 2, 6, 5 + prompt, gen = 2, 3 + seq_len = prompt + gen + m = _make_cba_metadata(max_batch=1, + K=K, + attn_len=10, + snap_len=6, + seq_len=seq_len, + prompt_len=prompt, + end_id=end_id, + batch=1) + for b in range(K): + m.cache_indirection[0, b, :] = b + m.cum_log_probs[0] = torch.tensor([-5.0, -1.0]) + # window rows distinguishable per beam + stop_window = torch.zeros((3, 1, K), dtype=torch.int32) + stop_window[:, 0, 0] = 100 + stop_window[:, 0, 1] = 200 + m.stop_past_tokens = stop_window + + # beam1 dominates: both slots descend from beam 1 + logits = torch.full((K, vocab), -50.0) + logits[1, 2] = 10.0 + logits[1, 3] = 9.0 + + beam_search_sampling_batch_cba( + logits=logits, + beam_width_in=K, + beam_width_out=K, + beam_search_args=m, + temperature=1.0, + early_stopping=0, + return_probs=False, + ) + assert m.predecessor_beams[0].tolist() == [1, 1] + # both window rows must now hold beam 1's history + assert stop_window[:, 0, 0].tolist() == [200, 200, 200] + assert stop_window[:, 0, 1].tolist() == [200, 200, 200] + + +@_kernel_test +def test_beam_search_sampling_batch_reorders_stop_window(): + """The ES=1 path must also reorder the stop-word window on beam swaps.""" + batch_size = 1 + beam_width = 2 + vocab_size = 6 + max_batch_size = 1 + seq_len = 6 + + seq_slots = torch.arange(batch_size, dtype=torch.int64) + stop_window = torch.zeros((3, max_batch_size, beam_width), + dtype=torch.int32) + stop_window[:, 0, 0] = 100 + stop_window[:, 0, 1] = 200 + metadata = BeamSearchMetadata( + cache_indirection=torch.zeros((max_batch_size, beam_width, seq_len + 1), + dtype=torch.int32), + cache_indirection_buffer=torch.full( + (max_batch_size, beam_width, seq_len + 1), -1, dtype=torch.int32), + cum_log_probs=torch.zeros((max_batch_size, beam_width), + dtype=torch.float32), + seq_slots=seq_slots, + seq_lens=torch.full((batch_size, ), seq_len, dtype=torch.int32), + finished_beams=torch.zeros((max_batch_size, beam_width), + dtype=torch.int32), + new_log_probs=torch.zeros((max_batch_size, beam_width), + dtype=torch.float32), + predecessor_beams=torch.zeros((max_batch_size, beam_width), + dtype=torch.int32), + seq_offsets=torch.arange(max_batch_size, dtype=torch.int64) * + beam_width, + beam_idx_arange=torch.arange(beam_width, dtype=torch.int32), + beam_gen_lengths=torch.zeros((max_batch_size, beam_width), + dtype=torch.int32), + stop_past_tokens=stop_window, + ) + metadata.cum_log_probs[0] = torch.tensor([-5.0, -1.0]) + + # beam 1 dominates: both slots descend from beam 1 + logits = torch.full((batch_size * beam_width, vocab_size), -50.0) + logits[1, 2] = 10.0 + logits[1, 3] = 9.0 + + beam_search_sampling_batch( + logits=logits, + beam_width_in=beam_width, + beam_width_out=beam_width, + beam_search_args=metadata, + temperature=1.0, + return_probs=False, + ) + assert metadata.predecessor_beams[0].tolist() == [1, 1] + assert stop_window[:, 0, 0].tolist() == [200, 200, 200] + assert stop_window[:, 0, 1].tolist() == [200, 200, 200] + + +@_kernel_test def test_beam_search_sampling_batch_disagg_handoff(): """Test context-first disagg beam handoff seeds gen-side beam scores.""" @@ -851,6 +1550,8 @@ def make_metadata(seq_len: int) -> BeamSearchMetadata: dtype=torch.int32), seq_offsets=seq_offsets, beam_idx_arange=beam_idx_arange, + beam_gen_lengths=torch.zeros((max_batch_size, beam_width), + dtype=torch.int32), ) torch.manual_seed(43) @@ -913,6 +1614,7 @@ def make_metadata(seq_len: int) -> BeamSearchMetadata: original_tokens=original_tokens, cache_indirection=disagg_metadata.cache_indirection, cum_log_probs=disagg_metadata.cum_log_probs, + beam_gen_lengths=disagg_metadata.beam_gen_lengths, ))))) for req_idx, seq_slot in enumerate(seq_slots.tolist()): first_gen_tokens = first_tokens[req_idx, :beam_width].tolist() @@ -1144,7 +1846,7 @@ class UutResultWrapper: # test def _uut(res=res): res.result = UutResult( - beam_history_builder=sampler._prepare_beam_history( + beam_history_builder=sampler._beam_search._prepare_beam_history( request, finish_reasons=torch.ones((beam_width, ), dtype=torch.int), d2h_copier=sampler._copy_to_host, @@ -1250,7 +1952,7 @@ def _uut(): cum_logprobs=cum_logprobs[batch_idx, :beam_width]) request.py_return_log_probs = False - sampler._finalize_beam(request, beam_history) + _finalize_beam(request, beam_history) token_history.append(deepcopy(request.get_tokens())) @@ -1285,6 +1987,92 @@ def _uut(): run_test_with_warmup(_uut_provider, max_sync_s=1) +def _beam_strategy(*, + early_stopping: BeamSearchEarlyStop, + beam_width_in: int = 2, + beam_width_out: int = 2, + temperature: float = 1.0, + length_penalty: float = 0.0, + diversity_rate: float = 0.0) -> BeamSearch: + return BeamSearch(tag="beam_search", + beam_width_in=beam_width_in, + beam_width_out=beam_width_out, + temperature=temperature, + length_penalty=length_penalty, + diversity_rate=diversity_rate, + early_stopping=early_stopping) + + +class TestBeamSearchStepFromStrategies: + """Cover the strategy-group -> BeamSearchStep construction. + + BeamSearchStep is abstract (`_select_and_update`), so only the concrete + subclasses may be instantiated. These tests pin that contract: they fail + with `TypeError: Can't instantiate abstract class BeamSearchStep` if the + extraction of the shared fields is ever done by constructing the base. + """ + + @staticmethod + def test_base_class_is_abstract(): + assert "_select_and_update" in _StrategyImpls.BeamSearchStep.__abstractmethods__ + with pytest.raises(TypeError, match="abstract"): + # Deliberately instantiating the abstract base: that is what this + # test pins, so mypy's (correct) complaint is expected here. + _StrategyImpls.BeamSearchStep( # type: ignore[abstract] + 2, 2, torch.ones(1), None, None) + + @pytest.mark.parametrize( + "early_stopping, expected_cls", + [ + (BeamSearchEarlyStop.TRUE, _StrategyImpls.RegularBeamSearchStep), + (BeamSearchEarlyStop.FALSE, _StrategyImpls.CBABeamSearchStep), + (BeamSearchEarlyStop.NEVER, _StrategyImpls.CBABeamSearchStep), + ], + ) + @staticmethod + @_kernel_test + def test_from_strategies_builds_concrete_impl(early_stopping, expected_cls): + strategies = [_beam_strategy(early_stopping=early_stopping)] + impl = expected_cls.from_strategies(strategies, + cuda_device=torch.device("cuda")) + assert type(impl) is expected_cls + assert impl._beam_width_in == 2 + assert impl._beam_width_out == 2 + # Unset length_penalty / diversity_rate stay None so the ops can skip + # the corresponding work. + assert impl._length_penalty is None + assert impl._diversity_rate is None + + @staticmethod + @_kernel_test + def test_from_strategies_carries_optional_fields(): + strategies = [ + _beam_strategy(early_stopping=BeamSearchEarlyStop.NEVER, + length_penalty=1.5, + diversity_rate=0.5) + ] + impl = _StrategyImpls.CBABeamSearchStep.from_strategies( + strategies, cuda_device=torch.device("cuda")) + assert impl._length_penalty is not None + assert impl._diversity_rate is not None + torch.testing.assert_close(impl._length_penalty, + torch.tensor([1.5], device="cuda")) + torch.testing.assert_close(impl._diversity_rate, + torch.tensor([0.5], device="cuda")) + assert impl._early_stopping is BeamSearchEarlyStop.NEVER + + @staticmethod + @_kernel_test + def test_common_fields_does_not_instantiate_base(): + """The shared extraction must return plain fields, not an instance.""" + strategies = [_beam_strategy(early_stopping=BeamSearchEarlyStop.FALSE)] + fields = _StrategyImpls.BeamSearchStep._common_fields( + strategies, torch.device("cuda")) + assert not isinstance(fields, _StrategyImpls.BeamSearchStep) + assert fields.beam_width_in == 2 + assert fields.beam_width_out == 2 + + @force_ampere # Save H100 resource class TestParameterValidation: """Ensure that unsupported request parameters do not crash/hang the engine.""" @@ -1340,56 +2128,38 @@ def _check_engine_responds(self, llm: LLM, input_prompts: list[str], @pytest.mark.timeout(120) @pytest.mark.threadleak(enabled=False) - def test_use_beam_search_false( + @pytest.mark.parametrize("use_beam_search", [False, None]) + def test_use_beam_search_disabled_rejects_multiple_returns( self, llm: LLM, input_prompts: list[str], fixed_params: dict[str, Any], batch_size: int, sampler_type: str, + use_beam_search: bool | None, ): + # best_of > 1 without beam search is greedy multi-return, which the LLM + # API rejects. Covers use_beam_search both explicitly False and omitted. if batch_size == 1: pytest.skip("Test does not depend on batch size") if sampler_type == "TorchSampler": pytest.skip("Test does not depend on sampler_type") assert fixed_params["max_beam_width"] > 2 + params = dict( + max_tokens=fixed_params["max_tokens"], + n=1, + best_of=fixed_params["max_beam_width"], + end_id=-1, + ) + if use_beam_search is not None: + params["use_beam_search"] = use_beam_search with pytest.raises( ValueError, match= ".*Greedy decoding in the LLM API does not allow multiple returns.*" ): _ = llm.generate(input_prompts, - sampling_params=SamplingParams( - max_tokens=fixed_params["max_tokens"], - n=1, - best_of=fixed_params["max_beam_width"], - use_beam_search=False, - end_id=-1, - )) - self._check_engine_responds(llm, input_prompts, fixed_params) - - @pytest.mark.timeout(120) - @pytest.mark.threadleak(enabled=False) - def test_use_beam_search_ommitted(self, llm: LLM, input_prompts: list[str], - fixed_params: dict[str, Any], - batch_size: int, sampler_type: str): - if batch_size == 1: - pytest.skip("Test does not depend on batch size") - if sampler_type == "TorchSampler": - pytest.skip("Test does not depend on sampler_type") - assert fixed_params["max_beam_width"] > 2 - with pytest.raises( - ValueError, - match= - ".*Greedy decoding in the LLM API does not allow multiple returns.*" - ): - _ = llm.generate(input_prompts, - sampling_params=SamplingParams( - max_tokens=fixed_params["max_tokens"], - n=1, - best_of=fixed_params["max_beam_width"], - end_id=-1, - )) + sampling_params=SamplingParams(**params)) self._check_engine_responds(llm, input_prompts, fixed_params) @pytest.mark.timeout(120) @@ -1404,17 +2174,32 @@ def test_smaller_beam_width( ): if batch_size == 1: pytest.skip("Test does not depend on batch size") - if sampler_type == "TorchSampler": - pytest.skip("Test does not depend on sampler_type") assert fixed_params["max_beam_width"] > 2 - with pytest.raises( - RequestError, - match=".*Request beam width 2 is not equal to max_beam_width 4.*" - ): + + # A beam width above max_beam_width is rejected: buffers are only + # allocated up to max_beam_width. + with pytest.raises(RequestError, match=".*exceeds max_beam_width.*"): _ = llm.generate(input_prompts, sampling_params=SamplingParams( max_tokens=fixed_params["max_tokens"], - n=1, + n=fixed_params["max_beam_width"] + 1, + best_of=fixed_params["max_beam_width"] + 1, + use_beam_search=True, + end_id=-1, + )) + self._check_engine_responds(llm, input_prompts, fixed_params) + + # A beam width below max_beam_width is rejected as well. TorchSampler + # can sample it, but 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; see TRTLLM-14792. + with pytest.raises(RequestError, + match=".*is not equal to max_beam_width.*"): + _ = llm.generate(input_prompts, + sampling_params=SamplingParams( + max_tokens=fixed_params["max_tokens"], + n=2, best_of=2, use_beam_search=True, end_id=-1, diff --git a/tests/unittest/_torch/sampler/test_beam_search_speculative_d2h.py b/tests/unittest/_torch/sampler/test_beam_search_speculative_d2h.py index c95ea62f9ef7..1d899c87f7d8 100644 --- a/tests/unittest/_torch/sampler/test_beam_search_speculative_d2h.py +++ b/tests/unittest/_torch/sampler/test_beam_search_speculative_d2h.py @@ -19,7 +19,7 @@ * parity vs the default (synchronous) beam-history D2H path, * the predictor-miss fallback in - `TorchSampler._prepare_beam_history._builder` + `BeamSearchHandler._prepare_beam_history._builder` (the synchronous `.cpu()` issued when the host-side predictor decided the step is non-terminal but the beam still finalizes), * the predictor-hit path that routes copies through the side stream. @@ -43,6 +43,7 @@ from tensorrt_llm import LLM, SamplingParams from tensorrt_llm._torch.models.checkpoints import HfCheckpointLoader from tensorrt_llm._torch.pyexecutor.sampler import SampleStateTorch, TorchSampler +from tensorrt_llm._torch.pyexecutor.sampler.beam_search import BeamSearchHandler from tensorrt_llm.executor.result import GenerationResult from tensorrt_llm.llmapi import KvCacheConfig @@ -157,7 +158,7 @@ def _run_with_env( Opt-in is via `TorchLlmArgs.enable_speculative_beam_history_d2h`. `predictor_override` patches - `TorchSampler._predict_beam_search_is_likely_finishing` and + `BeamSearchHandler.predict_is_likely_finishing` and `sampler_method_patches` patches arbitrary `TorchSampler` methods; either forces `TLLM_WORKER_USE_SINGLE_PROCESS=1` so class-level patches reach the sampler. `sampler_force_async_worker` enables the @@ -169,9 +170,7 @@ def _run_with_env( # sampler to run in-process so the patch is observed. monkeypatch.setenv("TLLM_WORKER_USE_SINGLE_PROCESS", "1") if predictor_override is not None: - monkeypatch.setattr( - TorchSampler, "_predict_beam_search_is_likely_finishing", predictor_override - ) + monkeypatch.setattr(BeamSearchHandler, "predict_is_likely_finishing", predictor_override) gc.collect(2) llm = _build_llm( diff --git a/tests/unittest/api_stability/references/trtllm_serve_api.yaml b/tests/unittest/api_stability/references/trtllm_serve_api.yaml index 99d263914f05..d9d07edd9e5e 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_api.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_api.yaml @@ -207,14 +207,14 @@ models: required: false length_penalty: kind: extension - type: float - default: 1.0 + type: Optional[float] + default: null status: stable required: false early_stopping: kind: extension - type: bool - default: false + type: Optional[int] + default: null status: stable required: false stop_token_ids: @@ -678,14 +678,14 @@ models: required: false length_penalty: kind: extension - type: float - default: 1.0 + type: Optional[float] + default: null status: stable required: false early_stopping: kind: extension - type: bool - default: false + type: Optional[int] + default: null status: stable required: false stop_token_ids: From 41a36a0b772f65026e7a23d319ad2a2059762666 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Fri, 31 Jul 2026 00:13:35 -0700 Subject: [PATCH 02/36] [TRTLLM-13234][fix] Match the actual beam-width rejection message in 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 --- tests/unittest/_torch/sampler/test_beam_search.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index e45f330cae7e..9700548cb758 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -2178,7 +2178,8 @@ def test_smaller_beam_width( # A beam width above max_beam_width is rejected: buffers are only # allocated up to max_beam_width. - with pytest.raises(RequestError, match=".*exceeds max_beam_width.*"): + with pytest.raises(RequestError, + match=".*is not equal to max_beam_width.*"): _ = llm.generate(input_prompts, sampling_params=SamplingParams( max_tokens=fixed_params["max_tokens"], From 621580cefabbcfc7eb58f680d8cb9c565d5589fa Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Mon, 3 Aug 2026 00:18:52 -0700 Subject: [PATCH 03/36] [TRTLLM-13234][fix] Address review: disagg check, early_stopping schema, 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 --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 5 +- tensorrt_llm/serve/openai_protocol.py | 27 +- .../_torch/sampler/test_beam_search.py | 330 ++++++++++++++++++ .../references/trtllm_serve_api.yaml | 4 +- 4 files changed, 359 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 3ba1777e5937..4f0dddc9721a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -4998,8 +4998,11 @@ def _validate_request(self, request: LlmRequest): # would be silently dropped -- reachable under aggregation but not # under disaggregation. Reject the combination until the handoff # carries the pool; TRTLLM-14792. + # NB: is_context_only_request is a property, but + # is_generation_only_request is a plain method -- it must be called, + # otherwise the bound method is truthy and this matches every request. if (request.is_context_only_request - or request.is_generation_only_request): + or request.is_generation_only_request()): early_stopping = _unwrap_singleton( sampling_config.early_stopping) if (early_stopping is not None diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 52fe584350e5..bda9c0e943cc 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -50,6 +50,25 @@ _LOGIT_BIAS_MIN = -100.0 _LOGIT_BIAS_MAX = 100.0 +# Beam-search stopping mode as exposed over HTTP, mirroring the HuggingFace +# Transformers interface. The engine encodes it as an integer, but that encoding +# is an implementation detail and is kept out of the public schema. +EarlyStopping: TypeAlias = Union[bool, Literal["never"]] + + +def _early_stopping_to_int(value: Optional[EarlyStopping]) -> Optional[int]: + """Translate the HF-style tri-state into the engine's integer encoding. + + Mirrors ``BeamSearchEarlyStop``: ``False`` -> 0, ``True`` -> 1, + ``"never"`` -> 2. ``None`` stays unset so the engine picks its default. + """ + if value is None: + return None + if value == "never": + return 2 + # NB: bool is a subclass of int, so this yields 1/0 for True/False. + return int(value) + def ensure_request_chat_template_allowed(request: Any, allow_request_chat_template: bool): @@ -535,7 +554,7 @@ class CompletionRequest(OpenAIBaseModel): # here would put every served request on the length-normalized, exhaustive # beam-search path. length_penalty: Optional[float] = None - early_stopping: Optional[int] = None + early_stopping: Optional[EarlyStopping] = None stop_token_ids: Optional[List[int]] = Field(default_factory=list) include_stop_str_in_output: bool = False ignore_eos: bool = False @@ -620,7 +639,7 @@ def to_sampling_params(self, min_p=self.min_p, repetition_penalty=self.repetition_penalty, length_penalty=self.length_penalty, - early_stopping=self.early_stopping, + early_stopping=_early_stopping_to_int(self.early_stopping), stop_token_ids=self.stop_token_ids, include_stop_str_in_output=self.include_stop_str_in_output, ignore_eos=self.ignore_eos, @@ -895,7 +914,7 @@ class ChatCompletionRequest(OpenAIBaseModel): # here would put every served request on the length-normalized, exhaustive # beam-search path. length_penalty: Optional[float] = None - early_stopping: Optional[int] = None + early_stopping: Optional[EarlyStopping] = None stop_token_ids: Optional[List[int]] = Field(default_factory=list) include_stop_str_in_output: bool = False ignore_eos: bool = False @@ -1040,7 +1059,7 @@ def to_sampling_params(self, min_p=self.min_p, repetition_penalty=self.repetition_penalty, length_penalty=self.length_penalty, - early_stopping=self.early_stopping, + early_stopping=_early_stopping_to_int(self.early_stopping), stop_token_ids=self.stop_token_ids, include_stop_str_in_output=self.include_stop_str_in_output, ignore_eos=self.ignore_eos, diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 9700548cb758..3867dcf549aa 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -44,6 +44,8 @@ BEAM_SEARCH_PAD_TOKEN, BeamSearch, BeamSearchEarlyStop, BeamSearchMetadata, CBAState, _StrategyImpls, beam_search_sampling_batch) from tensorrt_llm.bindings.executor import FinishReason +from tensorrt_llm.bindings.internal.batch_manager import \ + LlmRequest as CppLlmRequest from tensorrt_llm.executor import RequestError from tensorrt_llm.executor.result import (CompletionOutput, GenerationResult, Logprob) @@ -638,6 +640,79 @@ def test_beam_search_large_beam_width_regression( f"all {beam_width} beams are identical: {beam_sequences[0]}") +@pytest.mark.threadleak(enabled=False) +def test_beam_search_vbws_e2e(monkeypatch: pytest.MonkeyPatch, ) -> None: + """Variable-Beam-Width-Search through the full engine path. + + Drives beam_width_array end to end (scheduler -> ModelEngine -> + TorchSampler), which the operator-level tests cannot cover: the width + comes from get_beam_width_by_iter(), and that only advances with the + real decoding loop. + + NB: single request on purpose. ModelEngine requires every generation + request in a batch to report the same per-iteration beam width, and + beam_width_array is indexed by each request's own decoding_iter, so + several requests admitted at different times would desynchronize and + abort the batch (TRTLLM-14792). + """ + max_beam_width = 4 + beam_width_array = [2, 3, 4] + input_prompts = [[1, 2, 3]] + vocab_size = DummyConfig().vocab_size + # Two steps past the end of beam_width_array, so the width has to hold at + # its last entry (the clamp in get_beam_width_by_iter is exercised) while + # staying within what this path reliably completes. + # + # NB: capped deliberately. With max_tokens >= 6 this test intermittently + # fails to terminate, including with a constant beam_width_array (e.g. + # [4, 4, 4]), which suggests the cause is not width variation and may not + # be specific to VBWS at all. Not yet diagnosed; tracked separately rather + # than blocking the coverage this test does provide. + max_tokens = 5 + + checkpoint_loader = HfCheckpointLoader( + weight_loader=DummyWeightLoader(), + config_loader=DummyConfigLoader(), + ) + + gc.collect(2) # force destruction of any other LLM instances + with _single_process_context(): + llm = LLM( + model=_pl.Path("dummy_path"), + checkpoint_loader=checkpoint_loader, + sampler_type="TorchSampler", + max_beam_width=max_beam_width, + max_batch_size=max_beam_width, + max_seq_len=64, + kv_cache_config=KvCacheConfig(max_tokens=10000), + disable_overlap_scheduler=True, + cuda_graph_config=None, + ) + with llm: + sampling_params = SamplingParams( + max_tokens=max_tokens, + n=max_beam_width, + best_of=max_beam_width, + use_beam_search=True, + beam_width_array=beam_width_array, + end_id=-1, + ) + outputs = llm.generate(deepcopy(input_prompts), + sampling_params=deepcopy(sampling_params)) + + assert isinstance(outputs, list) + assert len(outputs) == len(input_prompts) + beams = outputs[0].outputs + assert len(beams) == max_beam_width, ( + f"expected {max_beam_width} beams, but got {len(beams)}") + for beam_idx, beam in enumerate(beams): + token_ids = beam.token_ids + assert token_ids is not None, f"beam {beam_idx} has no token_ids" + assert len(token_ids) > 0, f"beam {beam_idx} is empty" + assert all(0 <= t < vocab_size for t in token_ids), ( + f"beam {beam_idx} has out-of-vocab tokens: {token_ids}") + + ########################################################################### # Unit tests ########################################################################### @@ -1067,6 +1142,107 @@ def test_beam_topk_flashinfer_parity(): @_kernel_test +@_kernel_test +@pytest.mark.parametrize("length_penalty", [0.0, 1.0]) +def test_vbws_width_transition_with_length_penalty(length_penalty: float): + """A VBWS width change must keep per-beam lengths aligned to their beams. + + Widening from 2 to 3 beams reorders and duplicates source beams; the + per-beam generated lengths that length_penalty normalizes by must follow + the same permutation. Beams here have deliberately unequal lengths, so a + per-request (rather than per-beam) length would score them identically. + """ + batch_size = 1 + beam_width_in, beam_width_out = 2, 3 + vocab_size = 8 + max_batch_size = 2 + max_beam_width = 4 + seq_len = 6 + device = torch.device("cuda") + + seq_slots = torch.arange(batch_size, dtype=torch.int64, device=device) + slot = int(seq_slots[0]) + + cum_log_probs = torch.zeros((max_batch_size, max_beam_width), + dtype=torch.float32, + device=device) + # Equal cumulative scores, unequal lengths: with length_penalty > 0 the + # longer beam normalizes to a better score, so ranking must change. + cum_log_probs[slot, :beam_width_in] = torch.tensor([-2.0, -2.0], + device=device) + beam_gen_lengths = torch.zeros((max_batch_size, max_beam_width), + dtype=torch.int32, + device=device) + beam_gen_lengths[slot, :beam_width_in] = torch.tensor([2, 5], + dtype=torch.int32, + device=device) + + metadata = BeamSearchMetadata( + cache_indirection=torch.zeros( + (max_batch_size, max_beam_width, seq_len + 1), + dtype=torch.int32, + device=device), + cache_indirection_buffer=torch.full( + (max_batch_size, max_beam_width, seq_len + 1), + -1, + dtype=torch.int32, + device=device), + cum_log_probs=cum_log_probs, + seq_slots=seq_slots, + seq_lens=torch.full((batch_size, ), + seq_len, + dtype=torch.int32, + device=device), + finished_beams=torch.zeros((max_batch_size, max_beam_width), + dtype=torch.int32, + device=device), + new_log_probs=torch.zeros((max_batch_size, max_beam_width), + dtype=torch.float32, + device=device), + predecessor_beams=torch.zeros((max_batch_size, max_beam_width), + dtype=torch.int32, + device=device), + seq_offsets=torch.zeros((batch_size + 1, ), + dtype=torch.int32, + device=device), + beam_idx_arange=torch.arange(max_beam_width, + dtype=torch.int32, + device=device), + beam_gen_lengths=beam_gen_lengths, + ) + + logits = torch.zeros((batch_size * beam_width_in, vocab_size), + dtype=torch.float32, + device=device) + next_tokens, _ = beam_search_sampling_batch( + logits=logits, + beam_width_in=beam_width_in, + beam_width_out=beam_width_out, + beam_search_args=metadata, + temperature=1.0, + length_penalty=length_penalty, + return_probs=False, + ) + + # Rows are laid out at max_beam_width; only the first beam_width_out + # columns carry this step's beams. + assert next_tokens.shape[1] == max_beam_width + # The widened slots must be initialized, not left stale. + assert torch.isfinite(metadata.cum_log_probs[slot, :beam_width_out]).all() + # Slots beyond the current width stay untouched. + torch.testing.assert_close( + metadata.cum_log_probs[slot, beam_width_out:], + torch.zeros(max_beam_width - beam_width_out, device=device)) + + predecessors = metadata.predecessor_beams[slot, :beam_width_out] + assert int(predecessors.min()) >= 0 + assert int(predecessors.max()) < beam_width_in + if length_penalty > 0.0: + # Equal raw scores, so normalization by the per-beam length decides: + # the longer source beam (index 1, length 5) must win. + assert int(predecessors[0]) == 1 + + def test_beam_search_sampling_batch_diversity_rate(): """diversity_rate wired through beam_search_sampling_batch changes beam selection while stored cum_log_probs stay raw.""" @@ -1730,6 +1906,122 @@ def create_default_sampler(test_params: GeneralTestParams) -> TorchSampler: return sampler +def _vbws_request(beam_width_array: list[int] | None, + max_beam_width: int = 4) -> LlmRequest: + """A beam-search request carrying ``beam_width_array`` (VBWS).""" + sampling_params = SamplingParams(n=max_beam_width, + best_of=max_beam_width, + use_beam_search=True, + beam_width_array=beam_width_array) + return LlmRequest(request_id=0, + seq_slot=0, + max_new_tokens=16, + input_tokens=[1, 2, 3], + end_id=-1, + sampling_config=SamplingConfig( + sampling_params._get_sampling_config()), + is_streaming=False) + + +@pytest.mark.parametrize( + "beam_width_array, expected", + [ + # Widening: index is (iteration - 1), clamped at both ends. + ([2, 3, 4], [2, 2, 3, 4]), + # Narrowing: beams are dropped as decoding proceeds. + ([4, 3, 2], [4, 4, 3, 2]), + ], + ids=["widening", "narrowing"], +) +def test_vbws_beam_width_by_iter_follows_array(beam_width_array: list[int], + expected: list[int]): + """get_beam_width_by_iter walks beam_width_array as decoding advances. + + Covers both directions: widening adds beams, narrowing drops them. + """ + request = _vbws_request(beam_width_array) + actual = [] + for iteration in range(len(expected)): + request.decoding_iter = iteration + actual.append(request.get_beam_width_by_iter()) + assert actual == expected + + # for_next_iteration looks one step ahead, i.e. it is the same sequence + # shifted by one -- this is what feeds beam_width_out during sampling. + request.decoding_iter = 0 + assert request.get_beam_width_by_iter( + for_next_iteration=True) == expected[1] + + +def test_vbws_beam_width_by_iter_clamps_past_array_end(): + """Decoding longer than beam_width_array must hold the last width. + + Regression test for the C++ formula, which clamps the iteration index with + the global kMaxBeamWidthArrayLength (assuming a padded array) and therefore + reads past the end of the raw user array, returning garbage. The Python + override clamps with the actual array length instead; see + LlmRequest.get_beam_width_by_iter. + """ + beam_width_array = [2, 3, 4] + request = _vbws_request(beam_width_array) + # Run well past the end of the array. + for iteration in range(len(beam_width_array), len(beam_width_array) + 8): + request.decoding_iter = iteration + assert request.get_beam_width_by_iter() == beam_width_array[-1] + assert request.get_beam_width_by_iter( + for_next_iteration=True) == beam_width_array[-1] + + +def test_vbws_cpp_formula_reads_past_array_end(): + """Document why LlmRequest.get_beam_width_by_iter overrides the binding. + + The C++ implementation clamps the iteration index with the global + kMaxBeamWidthArrayLength rather than the actual array length, so once + decoding runs longer than the user's array it reads out of bounds and + returns arbitrary values (observed: 0, 32, 849 for a 3-entry array). + Values of 0 are particularly bad -- a zero beam width is not a valid + decoding state. TRTLLMSampler calls into C++ directly and is therefore + still affected; this test pins the divergence so the override is not + dropped as redundant. + """ + beam_width_array = [2, 3, 4] + request = _vbws_request(beam_width_array) + + # Within the array both agree. + for iteration in range(len(beam_width_array) + 1): + request.decoding_iter = iteration + assert (request.get_beam_width_by_iter() == + CppLlmRequest.get_beam_width_by_iter(request, False)) + + # Past the end the C++ formula diverges, and can even yield 0. + diverged = False + for iteration in range( + len(beam_width_array) + 1, + len(beam_width_array) + 8): + request.decoding_iter = iteration + assert request.get_beam_width_by_iter() == beam_width_array[-1] + if CppLlmRequest.get_beam_width_by_iter(request, + False) != beam_width_array[-1]: + diverged = True + assert diverged, ( + "C++ get_beam_width_by_iter no longer reads past the end of the " + "array; the Python override may now be redundant.") + + +def test_vbws_uniform_array_matches_fixed_width(): + """A constant beam_width_array must behave exactly like a fixed width.""" + max_beam_width = 4 + vbws = _vbws_request([max_beam_width] * 3, max_beam_width=max_beam_width) + fixed = _vbws_request(None, max_beam_width=max_beam_width) + for iteration in range(8): + vbws.decoding_iter = iteration + fixed.decoding_iter = iteration + assert vbws.get_beam_width_by_iter() == fixed.get_beam_width_by_iter() + assert vbws.get_beam_width_by_iter( + for_next_iteration=True) == fixed.get_beam_width_by_iter( + for_next_iteration=True) + + def test_create_beam_history(): """Test TorchSampler._create_beam_history method. @@ -2162,6 +2454,44 @@ def test_use_beam_search_disabled_rejects_multiple_returns( sampling_params=SamplingParams(**params)) self._check_engine_responds(llm, input_prompts, fixed_params) + @pytest.mark.timeout(120) + @pytest.mark.threadleak(enabled=False) + @pytest.mark.parametrize("early_stopping", [0, 2]) + def test_exhaustive_early_stopping_allowed_without_disagg( + self, + llm: LLM, + input_prompts: list[str], + fixed_params: dict[str, Any], + batch_size: int, + sampler_type: str, + early_stopping: int, + ): + # The exhaustive early_stopping modes are only rejected for + # disaggregated serving (the finished-candidate pool is not part of the + # handoff; TRTLLM-14792). Regular serving must still accept them. + # NB: guards against the disagg check matching every request, e.g. by + # testing a bound method rather than calling it. + if batch_size == 1: + pytest.skip("Test does not depend on batch size") + if sampler_type == "TRTLLMSampler": + pytest.skip("Exhaustive early_stopping check is TorchSampler-side") + outputs = llm.generate(input_prompts, + sampling_params=SamplingParams( + max_tokens=fixed_params["max_tokens"], + n=1, + best_of=fixed_params["max_beam_width"], + use_beam_search=True, + early_stopping=early_stopping, + end_id=-1, + )) + assert isinstance(outputs, list) + assert len(outputs) == len(input_prompts) + for output in outputs: + assert len(output.outputs) == 1 + token_ids = output.outputs[0].token_ids + assert token_ids is not None and len(token_ids) > 0 + self._check_engine_responds(llm, input_prompts, fixed_params) + @pytest.mark.timeout(120) @pytest.mark.threadleak(enabled=False) def test_smaller_beam_width( diff --git a/tests/unittest/api_stability/references/trtllm_serve_api.yaml b/tests/unittest/api_stability/references/trtllm_serve_api.yaml index d9d07edd9e5e..34f7e05bd495 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_api.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_api.yaml @@ -213,7 +213,7 @@ models: required: false early_stopping: kind: extension - type: Optional[int] + type: Optional[Union[bool, Literal['never']]] default: null status: stable required: false @@ -684,7 +684,7 @@ models: required: false early_stopping: kind: extension - type: Optional[int] + type: Optional[Union[bool, Literal['never']]] default: null status: stable required: false From a461406e9cbb22872d5ebd231b4b00bf634a07cc Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Mon, 3 Aug 2026 03:19:54 -0700 Subject: [PATCH 04/36] [TRTLLM-13234][test] Assert the per-iteration beam width in the VBWS 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 --- .../_torch/sampler/test_beam_search.py | 43 +++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 3867dcf549aa..fd9aeaeb807f 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -675,6 +675,27 @@ def test_beam_search_vbws_e2e(monkeypatch: pytest.MonkeyPatch, ) -> None: config_loader=DummyConfigLoader(), ) + # Record the width the engine actually uses at each decoding iteration. + # Asserting only on the outputs cannot distinguish a request that really + # walked [2, 3, 4] from one that ran at a constant width the whole time, + # so wrap the accessor every consumer (scheduler, ModelEngine, sampler) + # goes through and keep the current-iteration widths it hands out. + observed_widths: dict[int, int] = {} + unwrapped_get_beam_width_by_iter = LlmRequest.get_beam_width_by_iter + + def recording_get_beam_width_by_iter(self: LlmRequest, + for_next_iteration: bool = False + ) -> int: + width = unwrapped_get_beam_width_by_iter(self, for_next_iteration) + if not for_next_iteration: + # decoding_iter is 1-based once decoding starts; several callers + # ask per step, so keep the first answer for each iteration. + observed_widths.setdefault(self.decoding_iter, width) + return width + + monkeypatch.setattr(LlmRequest, "get_beam_width_by_iter", + recording_get_beam_width_by_iter) + gc.collect(2) # force destruction of any other LLM instances with _single_process_context(): llm = LLM( @@ -712,6 +733,28 @@ def test_beam_search_vbws_e2e(monkeypatch: pytest.MonkeyPatch, ) -> None: assert all(0 <= t < vocab_size for t in token_ids), ( f"beam {beam_idx} has out-of-vocab tokens: {token_ids}") + # The width must actually vary along beam_width_array and then hold at its + # last entry, rather than staying constant: decoding iteration i (1-based) + # uses beam_width_array[i - 1], clamped to the final entry once decoding + # outruns the array. + decoding_iters = sorted(it for it in observed_widths if it >= 1) + assert decoding_iters, ( + f"no decoding iterations recorded: {observed_widths}") + actual = [observed_widths[it] for it in decoding_iters] + expected = [ + beam_width_array[min(it, len(beam_width_array)) - 1] + for it in decoding_iters + ] + assert actual == expected, ( + f"beam width per decoding iteration {decoding_iters} was {actual}, " + f"expected {expected} from beam_width_array={beam_width_array}") + # Guard against the whole run happening at a single width, which the + # per-iteration comparison above would still accept if the engine only + # ever reported one iteration. + assert set(actual) == set(beam_width_array), ( + f"expected every width in {beam_width_array} to be exercised, " + f"but only saw {sorted(set(actual))}") + ########################################################################### # Unit tests From 08bcc638c187b57e636235a4eb0dd0c675f7badb Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Mon, 3 Aug 2026 08:15:25 -0700 Subject: [PATCH 05/36] [TRTLLM-13234][fix] Clamp getBeamWidthByIter with the actual beam width 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 --- cpp/tensorrt_llm/batch_manager/llmRequest.cpp | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/cpp/tensorrt_llm/batch_manager/llmRequest.cpp b/cpp/tensorrt_llm/batch_manager/llmRequest.cpp index d5466b4a7539..fa0f69c79235 100644 --- a/cpp/tensorrt_llm/batch_manager/llmRequest.cpp +++ b/cpp/tensorrt_llm/batch_manager/llmRequest.cpp @@ -37,13 +37,16 @@ runtime::SizeType32 GenericLlmRequest::getBeamWidthByIter(bool { runtime::SizeType32 beamWidth = mSamplingConfig.beamWidth; // For non-Variable-Beam-Width-Search auto const& beamWidthArray = mSamplingConfig.beamWidthArray; - if (beamWidthArray.has_value()) + if (beamWidthArray.has_value() && !beamWidthArray.value().empty() && !beamWidthArray.value()[0].empty()) { + auto const& requestBeamWidthArray = beamWidthArray.value()[0]; auto const iter = mDecodingIter + (forNextIteration ? 1 : 0); - // Clamped `decodingIter` into [0,kMaxBeamWidthArrayLength-1] as index - int const index - = std::max(std::min(iter, static_cast(tensorrt_llm::kernels::kMaxBeamWidthArrayLength)) - 1, 0); - beamWidth = beamWidthArray.value()[0][index]; + // Clamp `decodingIter` with the actual array length, so that decoding + // longer than the array holds the last width instead of reading past + // the end. kMaxBeamWidthArrayLength is only the capacity limit; the + // user array is not padded up to it. + int const index = std::max(std::min(iter, static_cast(requestBeamWidthArray.size())) - 1, 0); + beamWidth = requestBeamWidthArray[index]; } return beamWidth; } From 06135e813d012f010823fd5364ae274f30f31dcc Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Mon, 3 Aug 2026 18:47:06 -0700 Subject: [PATCH 06/36] [TRTLLM-13234][test] Decode further past the beam width array in the 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 --- .../_torch/sampler/test_beam_search.py | 21 +++++++++++-------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index fd9aeaeb807f..508724185a9f 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -659,16 +659,19 @@ def test_beam_search_vbws_e2e(monkeypatch: pytest.MonkeyPatch, ) -> None: beam_width_array = [2, 3, 4] input_prompts = [[1, 2, 3]] vocab_size = DummyConfig().vocab_size - # Two steps past the end of beam_width_array, so the width has to hold at - # its last entry (the clamp in get_beam_width_by_iter is exercised) while - # staying within what this path reliably completes. + # Decode well past the end of beam_width_array, so the width has to hold at + # its last entry for most of the run: that clamp is the interesting part, + # since getting it wrong reads past the array. # - # NB: capped deliberately. With max_tokens >= 6 this test intermittently - # fails to terminate, including with a constant beam_width_array (e.g. - # [4, 4, 4]), which suggests the cause is not width variation and may not - # be specific to VBWS at all. Not yet diagnosed; tracked separately rather - # than blocking the coverage this test does provide. - max_tokens = 5 + # NB: running beyond the array used to hang here. Both getBeamWidthByIter + # implementations clamp the iteration index, but the C++ one used the + # global kMaxBeamWidthArrayLength instead of the array's own length and + # returned an out-of-bounds width once decoding outran the array. The C++ + # micro-batch scheduler reads that width, so the request was never admitted + # into the generation batch again and decoding stalled forever. Exercising + # the clamp therefore requires the C++ fix in llmRequest.cpp -- against an + # older libtensorrt_llm.so this test hangs rather than fails. + max_tokens = 12 checkpoint_loader = HfCheckpointLoader( weight_loader=DummyWeightLoader(), From 181060fe9f835885c254644de98394403f2391d0 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 5 Aug 2026 06:18:15 -0700 Subject: [PATCH 07/36] [TRTLLM-13234][fix] Exclude every kind of dummy request from the beam 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 --- .../_torch/pyexecutor/model_engine.py | 12 ++-- .../_torch/sampler/test_beam_search.py | 64 +++++++++++-------- 2 files changed, 45 insertions(+), 31 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index dcc55f586efb..865bed3b5d4b 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -5248,13 +5248,13 @@ def append_cross_attention_state(request: LlmRequest, # variable-beam-width request narrows or widens per iteration, so # the widths can still diverge mid-batch. Compare the # *per-iteration* width: py_beam_width is fixed at admission and - # would be identical across those requests. CUDA-graph padding - # requests are built at the engine width and appended after - # scheduling, so they are excluded -- they carry no user request - # and would otherwise trip this on an ordinary padded batch. + # would be identical across those requests. Dummy requests are + # excluded -- they carry no user request and are built at their own + # width (CUDA-graph padding at the engine width, attention-DP and + # warmup dummies at width one), so they would otherwise trip this + # on an ordinary padded batch. real_requests = [ - req for req in generation_requests - if not req.is_cuda_graph_dummy + req for req in generation_requests if not req.is_dummy ] iter_widths = { req.get_beam_width_by_iter() diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 508724185a9f..5759fadb39a0 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -2018,40 +2018,33 @@ def test_vbws_beam_width_by_iter_clamps_past_array_end(): for_next_iteration=True) == beam_width_array[-1] -def test_vbws_cpp_formula_reads_past_array_end(): - """Document why LlmRequest.get_beam_width_by_iter overrides the binding. - - The C++ implementation clamps the iteration index with the global - kMaxBeamWidthArrayLength rather than the actual array length, so once - decoding runs longer than the user's array it reads out of bounds and - returns arbitrary values (observed: 0, 32, 849 for a 3-entry array). - Values of 0 are particularly bad -- a zero beam width is not a valid - decoding state. TRTLLMSampler calls into C++ directly and is therefore - still affected; this test pins the divergence so the override is not - dropped as redundant. +def test_vbws_cpp_formula_matches_past_array_end(): + """The C++ and Python clamps must agree once decoding outruns the array. + + The C++ implementation used to clamp the iteration index with the global + kMaxBeamWidthArrayLength rather than the actual array length, so it read + out of bounds and returned arbitrary widths (observed: 0, 32, 849 for a + 3-entry array). That starved the request in the C++ micro-batch scheduler + and hung decoding; it is fixed in llmRequest.cpp. TRTLLMSampler and the + scheduler call into C++ directly, so pin the agreement here -- a failure + means the two clamps have drifted apart again. """ beam_width_array = [2, 3, 4] request = _vbws_request(beam_width_array) - # Within the array both agree. - for iteration in range(len(beam_width_array) + 1): + for iteration in range(len(beam_width_array) + 8): request.decoding_iter = iteration assert (request.get_beam_width_by_iter() == CppLlmRequest.get_beam_width_by_iter(request, False)) + assert (request.get_beam_width_by_iter( + for_next_iteration=True) == CppLlmRequest.get_beam_width_by_iter( + request, True)) - # Past the end the C++ formula diverges, and can even yield 0. - diverged = False - for iteration in range( - len(beam_width_array) + 1, - len(beam_width_array) + 8): + # Past the end both must hold the last entry rather than read past it. + for iteration in range(len(beam_width_array), len(beam_width_array) + 8): request.decoding_iter = iteration - assert request.get_beam_width_by_iter() == beam_width_array[-1] - if CppLlmRequest.get_beam_width_by_iter(request, - False) != beam_width_array[-1]: - diverged = True - assert diverged, ( - "C++ get_beam_width_by_iter no longer reads past the end of the " - "array; the Python override may now be redundant.") + assert CppLlmRequest.get_beam_width_by_iter( + request, False) == beam_width_array[-1] def test_vbws_uniform_array_matches_fixed_width(): @@ -2068,6 +2061,27 @@ def test_vbws_uniform_array_matches_fixed_width(): for_next_iteration=True) +def test_vbws_dummy_requests_excluded_from_width_check(): + """Every kind of dummy must be excluded from the mixed-width guard. + + ModelEngine rejects a generation batch whose requests report different + per-iteration beam widths, but dummy requests carry no user request and are + built at their own width -- CUDA-graph padding at the engine width, + attention-DP and warmup dummies at width one. Filtering only CUDA-graph + dummies let an attention-DP dummy abort an otherwise valid beam-search + batch, so the guard filters on `is_dummy`; pin that it covers all three + flags. + """ + for flag in ("is_cuda_graph_dummy", "is_attention_dp_dummy", + "is_dummy_request"): + request = _vbws_request([2, 3, 4]) + assert not request.is_dummy, f"{flag}: unset request must not be dummy" + setattr(request, flag, True) + assert request.is_dummy, ( + f"{flag} is not covered by is_dummy, so such requests would reach " + "the mixed-beam-width check in ModelEngine and abort the batch") + + def test_create_beam_history(): """Test TorchSampler._create_beam_history method. From 2eb914d7df00eea6e5636cf893102fb1adedf488 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 5 Aug 2026 06:54:01 -0700 Subject: [PATCH 08/36] [TRTLLM-13234][fix] Offset generation logits by the static beam width 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 --- .../_torch/pyexecutor/sampler/beam_search.py | 27 ++++++++++++++++--- .../_torch/pyexecutor/sampler/sampler.py | 14 ++++++++-- .../pyexecutor/sampler/sampler_common.py | 7 +++++ .../pyexecutor/sampler/sampler_strategy.py | 21 +++++++++++++++ 4 files changed, 64 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py index 061861220e1a..aa9aaf9d0e0a 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py @@ -364,6 +364,7 @@ def _beam_step_preprocess( logits: torch.Tensor, *, beam_width_in: int, + row_stride: int | None = None, temperature: float | None, return_probs: bool, args: "BeamSearchMetadata", @@ -373,12 +374,28 @@ def _beam_step_preprocess( Applies temperature, snapshots the cache indirection into its buffer, and returns ``(logprobs, softmax, batch_size)``. ``softmax`` is None when ``return_probs`` is False. + + ``row_stride`` is how many rows the forward path allocated per request, + which is the static admission width. It differs from ``beam_width_in`` + under Variable-Beam-Width-Search, where only the first ``beam_width_in`` + rows of each request hold live beams and the rest are padding. Reshaping + by ``beam_width_in`` alone would silently mix rows across requests, since + ``view`` accepts any shape whose element count divides. Defaults to + ``beam_width_in`` for callers whose layout already matches. """ assert logits.dim() == 2, "logits should be 2D: [batch_size * beam_width, vocab_size]" - batch_size, vocab_size = logits.size() - batch_size = batch_size // beam_width_in + num_rows, vocab_size = logits.size() + if row_stride is None: + row_stride = beam_width_in + assert row_stride >= beam_width_in, ( + f"row_stride ({row_stride}) must cover beam_width_in ({beam_width_in})" + ) + batch_size = num_rows // row_stride - logits = logits.view(batch_size, beam_width_in, vocab_size) + logits = logits.view(batch_size, row_stride, vocab_size) + if row_stride != beam_width_in: + # Drop the padding rows; the live beams are the leading ones. + logits = logits[:, :beam_width_in, :] if temperature is not None and temperature != 0: logits = logits / max(temperature, 1e-5) softmax: Optional[torch.Tensor] = None @@ -508,6 +525,7 @@ def beam_search_sampling_batch( *, beam_width_in: int, beam_width_out: int, + row_stride: int | None = None, beam_search_args: BeamSearchMetadata, temperature: float | None, length_penalty: "torch.Tensor | float | None" = None, @@ -528,6 +546,7 @@ def beam_search_sampling_batch( logprobs, softmax, batch_size = _beam_step_preprocess( logits, beam_width_in=beam_width_in, + row_stride=row_stride, temperature=temperature, return_probs=return_probs, args=beam_search_args, @@ -866,6 +885,7 @@ def beam_search_sampling_batch_cba( *, beam_width_in: int, beam_width_out: int, + row_stride: int | None = None, beam_search_args: BeamSearchMetadata, temperature: float | None, early_stopping: int, # BeamSearchEarlyStop @@ -912,6 +932,7 @@ def beam_search_sampling_batch_cba( logprobs, softmax, batch_size = _beam_step_preprocess( logits, beam_width_in=beam_width_in, + row_stride=row_stride, temperature=temperature, return_probs=return_probs, args=args, diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 2617fa4bfab6..73db892f2499 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -3368,9 +3368,19 @@ def _select_generated_logits( req_num_generation_steps_list, dtype=torch.int32, pin_memory=prefer_pinned() ) - # context requests do not have multiple beams yet, so beam width may differ in mixed batches + # Rows in the logits tensor, i.e. how many rows each request occupies. + # ModelEngine lays out generation requests at the *static* admission + # width (py_beam_width), not the per-iteration width: with a variable + # beam width array the two differ, and offsetting by the narrower + # per-iteration width would make every request after the first read + # another request's rows. logits.view() succeeds for any shape whose + # element count divides, so that is silent corruption rather than an + # error. Match the layout here and slice down to the per-iteration + # width where the beams are actually consumed. TRTLLMSampler already + # offsets by the static width for the same reason. + # NB: context requests do not have multiple beams yet, hence the 1s. req_num_beams_list = [1] * len(finished_context_requests) + [ - req.get_beam_width_by_iter(False) for req in scheduled_requests.generation_requests + req.py_beam_width for req in scheduled_requests.generation_requests ] req_num_beams = torch.tensor( req_num_beams_list, dtype=torch.int32, pin_memory=prefer_pinned() diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_common.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_common.py index 0a934587a435..e787d69f970d 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_common.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_common.py @@ -76,6 +76,9 @@ class UtilsSamplingParams: min_p: Optional[float] = None beam_width_in: Optional[int] = None beam_width_out: Optional[int] = None + # Rows the forward path allocated per request (static admission width). + # Equals beam_width_in unless the request uses a variable beam width array. + row_stride: Optional[int] = None top_p_decay: Optional[float] = None top_p_min: Optional[float] = None top_p_reset_ids: Optional[int] = None @@ -169,6 +172,9 @@ def _request_get_sampling_params(request: LlmRequest) -> UtilsSamplingParams: top_p_reset_ids = _unwrap_singleton(cast(Optional[list[int]], sampling_config.top_p_reset_ids)) beam_width_out = _get_beam_width_out(request) beam_width_in = _get_beam_width_in(request) + # ModelEngine lays generation rows out at the static admission width; see + # the row_stride note in _beam_step_preprocess. + row_stride = 1 if request.is_context_init_state else request.py_beam_width use_beam_search = _get_max_beam_width(request) > 1 length_penalty = _unwrap_singleton(cast(Optional[list[float]], sampling_config.length_penalty)) beam_search_diversity_rate = _unwrap_singleton( @@ -183,6 +189,7 @@ def _request_get_sampling_params(request: LlmRequest) -> UtilsSamplingParams: min_p=min_p, beam_width_in=beam_width_in, beam_width_out=beam_width_out, + row_stride=row_stride, use_beam_search=use_beam_search, top_p_decay=top_p_decay, top_p_min=top_p_min, diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py index 4cba16d8754a..ab381f8021ab 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py @@ -130,6 +130,9 @@ class BeamSearch(NamedTuple): length_penalty: float diversity_rate: float early_stopping: BeamSearchEarlyStop + # Appended last on purpose: _common_fields() reads the fields above by + # position, so inserting earlier would shift those indices. + row_stride: int = 0 GREEDY: Greedy = ("greedy", None) @@ -242,6 +245,7 @@ def resolve_sampling_strategy(params: UtilsSamplingParams, *, vocab_size: int) - length_penalty=params.length_penalty or 0.0, diversity_rate=params.beam_search_diversity_rate or 0.0, early_stopping=BeamSearchEarlyStop.from_raw(params.early_stopping), + row_stride=params.row_stride or params.beam_width_in, ) # NB: not greedy, hence top_p != 0 if specified @@ -332,7 +336,9 @@ def sample( length_penalty, beam_search_diversity_rate, early_stopping, + *_, ): + row_stride = cast(BeamSearch, strategy).row_stride assert group_metadata is not None and isinstance(group_metadata, BeamSearchMetadata), ( "BeamSearchMetadata is required for beam_search_sampling_batch" ) @@ -341,6 +347,7 @@ def sample( logits, beam_width_in=cast(int, beam_width_in), beam_width_out=cast(int, beam_width_out), + row_stride=cast(int, row_stride), beam_search_args=group_metadata, temperature=cast(float, temperature), early_stopping=cast(int, early_stopping), @@ -353,6 +360,7 @@ def sample( logits, beam_width_in=cast(int, beam_width_in), beam_width_out=cast(int, beam_width_out), + row_stride=cast(int, row_stride), beam_search_args=group_metadata, temperature=cast(float, temperature), length_penalty=cast(float, length_penalty), @@ -1026,6 +1034,7 @@ class CommonFields: beam_width_in: int beam_width_out: int + row_stride: int temperature: torch.Tensor length_penalty: Optional[torch.Tensor] diversity_rate: Optional[torch.Tensor] @@ -1034,6 +1043,7 @@ def __init__( self, beam_width_in: int, beam_width_out: int, + row_stride: int, temperature: torch.Tensor, length_penalty: Optional[torch.Tensor], diversity_rate: Optional[torch.Tensor], @@ -1042,6 +1052,7 @@ def __init__( ): self._beam_width_in = beam_width_in self._beam_width_out = beam_width_out + self._row_stride = row_stride self._temperature = temperature self._length_penalty = length_penalty self._diversity_rate = diversity_rate @@ -1072,6 +1083,9 @@ def _common_fields( narrowed_strats = cast(list[BeamSearch], strategies) (beam_width_in,) = set(strat[1] for strat in narrowed_strats) (beam_width_out,) = set(strat[2] for strat in narrowed_strats) + # Grouping already keys on the strategy tuple, so a group cannot + # mix row strides; assert rather than silently pick one. + (row_stride,) = set(strat.row_stride or beam_width_in for strat in narrowed_strats) temperature = _StrategyImpls.BeamSearchStep._make_tensor( [strat[3] or 1.0 for strat in narrowed_strats], torch.float32, cuda_device ) @@ -1090,6 +1104,7 @@ def _common_fields( return _StrategyImpls.BeamSearchStep.CommonFields( beam_width_in=beam_width_in, beam_width_out=beam_width_out, + row_stride=row_stride, temperature=temperature, length_penalty=length_penalty, diversity_rate=diversity_rate, @@ -1104,6 +1119,7 @@ def from_strategies( return cls( fields.beam_width_in, fields.beam_width_out, + fields.row_stride, fields.temperature, fields.length_penalty, fields.diversity_rate, @@ -1141,6 +1157,7 @@ def _select_and_update( logits, beam_width_in=self._beam_width_in, beam_width_out=self._beam_width_out, + row_stride=self._row_stride, beam_search_args=group_metadata, temperature=None, length_penalty=self._length_penalty, @@ -1155,6 +1172,7 @@ def __init__( self, beam_width_in: int, beam_width_out: int, + row_stride: int, temperature: torch.Tensor, length_penalty: Optional[torch.Tensor], diversity_rate: Optional[torch.Tensor], @@ -1165,6 +1183,7 @@ def __init__( super().__init__( beam_width_in, beam_width_out, + row_stride, temperature, length_penalty, diversity_rate, @@ -1184,6 +1203,7 @@ def from_strategies( return cls( fields.beam_width_in, fields.beam_width_out, + fields.row_stride, fields.temperature, fields.length_penalty, fields.diversity_rate, @@ -1198,6 +1218,7 @@ def _select_and_update( logits, beam_width_in=self._beam_width_in, beam_width_out=self._beam_width_out, + row_stride=self._row_stride, beam_search_args=group_metadata, temperature=None, early_stopping=self._early_stopping, From 196365e9af59341d02b7e8e28b1ceffd3ef9e494 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 5 Aug 2026 07:04:16 -0700 Subject: [PATCH 09/36] [TRTLLM-13234][fix] Pass row_stride in the abstract BeamSearchStep test 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 --- tests/unittest/_torch/sampler/test_beam_search.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 5759fadb39a0..b3513536f04f 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -2371,7 +2371,7 @@ def test_base_class_is_abstract(): # Deliberately instantiating the abstract base: that is what this # test pins, so mypy's (correct) complaint is expected here. _StrategyImpls.BeamSearchStep( # type: ignore[abstract] - 2, 2, torch.ones(1), None, None) + 2, 2, 2, torch.ones(1), None, None) @pytest.mark.parametrize( "early_stopping, expected_cls", From a4bfc227e13ae9529549c191d6e41bc6d009a4ee Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 5 Aug 2026 08:17:40 -0700 Subject: [PATCH 10/36] [TRTLLM-13234][fix] Serve every early_stopping mode from the CBA path 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 --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 60 ++++-- .../_torch/pyexecutor/sampler/beam_search.py | 39 ++-- .../pyexecutor/sampler/sampler_strategy.py | 62 +++--- .../_torch/sampler/test_beam_search.py | 178 +++++++----------- 4 files changed, 156 insertions(+), 183 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 4f0dddc9721a..2fd71afded98 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -87,8 +87,6 @@ ResourceManagerType, request_context) from .sampler import (AsyncWorkerMixin, Sampler, SamplerEvent, SampleState, SampleStateTensors, TRTLLMSampler) -from .sampler.beam_search import BeamSearchEarlyStop -from .sampler.sampler_common import _unwrap_singleton from .scheduler import (RequestScheduler, ScheduledRequests, SerializableSchedulerOutput, WaitingQueue, create_waiting_queue) @@ -4991,29 +4989,53 @@ def _validate_request(self, request: LlmRequest): f"is not equal to max_beam_width {self.max_beam_width}. " "This is not supported!") - # Exhaustive early_stopping keeps a pool of finished candidates, - # which the context server can already populate on its first step. - # That pool is not part of the disaggregated handoff and is cleared - # on the generation side, so a completion the context phase found - # would be silently dropped -- reachable under aggregation but not - # under disaggregation. Reject the combination until the handoff - # carries the pool; TRTLLM-14792. + # Variable-Beam-Width-Search is only defined for a non-decreasing + # array. The documented semantics (see getBeamWidthByIter in + # llmRequest.h) only cover widening, ending at the full width, and + # both samplers rely on that: the per-step ops write the leading + # beam_width_out rows while finalize reads py_beam_width (the + # array maximum) of them, so a narrowing array leaves the trailing + # rows holding ancestry from an earlier, wider step. The C++ + # finalize path has the same gap -- it indexes with nBeamWidth + # only. Reject rather than emit silently stale beams; narrowing + # can be allowed once its semantics are defined (TRTLLM-14792). + beam_width_array = sampling_config.beam_width_array + if beam_width_array: + if isinstance(beam_width_array[0], (list, tuple)): + beam_width_array = beam_width_array[0] + if any(b < a + for a, b in zip(beam_width_array, beam_width_array[1:])): + raise ValueError( + f"beam_width_array {list(beam_width_array)} decreases; " + "only non-decreasing arrays are supported for " + "Variable-Beam-Width-Search.") + + # Beam search keeps a pool of finished candidates, which the + # context server can already populate on its first step. That pool + # is not part of the disaggregated 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 -- reachable under aggregation but not under + # disaggregation. + # + # This applies to every early_stopping mode and to both samplers: + # TorchSampler now serves all modes from the candidate-beams-array + # path, and the C++ decoder behind TRTLLMSampler maintains the pool + # unconditionally (beamSearchLayer.cu assigns numBeamsCBA with no + # mode check). Reject until the handoff carries the pool; + # TRTLLM-14792. # NB: is_context_only_request is a property, but # is_generation_only_request is a plain method -- it must be called, # otherwise the bound method is truthy and this matches every request. if (request.is_context_only_request or request.is_generation_only_request()): - early_stopping = _unwrap_singleton( - sampling_config.early_stopping) - if (early_stopping is not None - and BeamSearchEarlyStop.from_raw(early_stopping) - is not BeamSearchEarlyStop.TRUE): + if sampling_config.beam_width > 1: raise ValueError( - f"Beam search early_stopping={early_stopping} is not " - "supported with disaggregated serving: the finished-" - "candidate pool is not transferred between the context " - "and generation servers. Use the default " - "(early_stopping=True).") + "Beam search is not supported with disaggregated " + "serving: the finished-candidate pool built during the " + "context phase is not transferred to the generation " + "server, so completions found there would be silently " + "dropped. Use beam_width=1.") # Check token ID ranges self._validate_token_id_range(request) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py index aa9aaf9d0e0a..f4b673660872 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py @@ -842,13 +842,22 @@ def _take(src: torch.Tensor) -> torch.Tensor: # decrease monotonically with sequence length, so longer sequences only get # less attractive). min_kept = top_normed.gather(1, (caps - 1).long()).view(-1) - if early_stopping != BeamSearchEarlyStop.FALSE: - max_gen = (max_seq_len - prompts.view(-1)).to(cand_len.dtype) - bound_len = torch.where(exponent.view(-1) > 0, max_gen, cand_len.view(-1)) + # `min_kept > neg_inf` means the pool holds `caps` finished hypotheses, + # i.e. the C++ `numBeamsCBA[slot] >= nBM` test. + pool_full = min_kept > neg_inf + if early_stopping == BeamSearchEarlyStop.TRUE: + # HF `True`: stop as soon as `beam_width` finished candidates exist, + # without weighing what is still attainable. Matches C++ + # beamStage3Kernel, which short-circuits to done here. + done = pool_full else: - bound_len = cand_len.view(-1) - best_attainable = cand_cum[:, 0] / bound_len.to(cand_cum.dtype).pow(exponent.view(-1)) - done = (min_kept > neg_inf) & (min_kept >= best_attainable) + if early_stopping != BeamSearchEarlyStop.FALSE: + max_gen = (max_seq_len - prompts.view(-1)).to(cand_len.dtype) + bound_len = torch.where(exponent.view(-1) > 0, max_gen, cand_len.view(-1)) + else: + bound_len = cand_len.view(-1) + best_attainable = cand_cum[:, 0] / bound_len.to(cand_cum.dtype).pow(exponent.view(-1)) + done = pool_full & (min_kept >= best_attainable) # Reorder the finish handler's rolling stop-word window to follow the # beam swap (matching stays correct across swaps). @@ -1166,12 +1175,18 @@ def _prepare_beam_search( def _request_uses_cba(request: LlmRequest) -> bool: - """Whether this beam-search request runs on the candidate-beams-array - path (exhaustive early_stopping modes).""" - early_stopping = _unwrap_singleton( - cast(Optional[list[int]], request.sampling_config.early_stopping) - ) - return BeamSearchEarlyStop.from_raw(early_stopping) != BeamSearchEarlyStop.TRUE + """Whether this beam-search request runs on the candidate-beams-array path. + + Every early_stopping mode does. The modes differ only in the done verdict + computed inside the CBA step (TRUE stops as soon as the pool is full; + FALSE / NEVER additionally weigh what is still attainable), matching the + C++ decoder, which maintains the pool for all modes. + + Kept as a predicate rather than inlined: it marks the places that depend on + the CBA state existing, and it is the hook to restore should a + pool-free path be reintroduced. + """ + return True def _prepare_beam_history_cba( diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py index ab381f8021ab..6bc94ad6929f 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py @@ -342,31 +342,22 @@ def sample( assert group_metadata is not None and isinstance(group_metadata, BeamSearchMetadata), ( "BeamSearchMetadata is required for beam_search_sampling_batch" ) - if cast(int, early_stopping) != BeamSearchEarlyStop.TRUE: - tokens, softmax = beam_search_sampling_batch_cba( - logits, - beam_width_in=cast(int, beam_width_in), - beam_width_out=cast(int, beam_width_out), - row_stride=cast(int, row_stride), - beam_search_args=group_metadata, - temperature=cast(float, temperature), - early_stopping=cast(int, early_stopping), - length_penalty=cast(float, length_penalty), - diversity_rate=cast(float, beam_search_diversity_rate), - return_probs=return_probs, - ) - else: - tokens, softmax = beam_search_sampling_batch( - logits, - beam_width_in=cast(int, beam_width_in), - beam_width_out=cast(int, beam_width_out), - row_stride=cast(int, row_stride), - beam_search_args=group_metadata, - temperature=cast(float, temperature), - length_penalty=cast(float, length_penalty), - diversity_rate=cast(float, beam_search_diversity_rate), - return_probs=return_probs, - ) + # Every early_stopping mode goes through the candidate-beams-array + # path: TRUE differs only in the done verdict (pool full, without + # weighing attainability), matching the C++ decoder, which keeps + # the pool for all modes. + tokens, softmax = beam_search_sampling_batch_cba( + logits, + beam_width_in=cast(int, beam_width_in), + beam_width_out=cast(int, beam_width_out), + row_stride=cast(int, row_stride), + beam_search_args=group_metadata, + temperature=cast(float, temperature), + early_stopping=cast(int, early_stopping), + length_penalty=cast(float, length_penalty), + diversity_rate=cast(float, beam_search_diversity_rate), + return_probs=return_probs, + ) return tokens, softmax, cast(float, temperature) @@ -1311,15 +1302,12 @@ def sample_grouped_strategies( strategy_impl_cls = _StrategyImpls.MinPWithProbs case "greedy": strategy_impl_cls = _StrategyImpls.GreedyWithProbs - case ("beam_search", beam_width_in_key, _, early_stopping_key): + case ("beam_search", beam_width_in_key, _, _): beam_width_in = beam_width_in_key - # Beam search encodes with-probs as a constructor flag, not a - # subclass; the stopping mode selects the class. - strategy_impl_cls = ( - _StrategyImpls.RegularBeamSearchStep - if early_stopping_key == BeamSearchEarlyStop.TRUE - else _StrategyImpls.CBABeamSearchStep - ) + # Beam search encodes with-probs as a constructor flag, not + # a subclass. Every stopping mode uses the CBA step; the + # mode only changes the done verdict inside it. + strategy_impl_cls = _StrategyImpls.CBABeamSearchStep case _: raise NotImplementedError("Unsupported strategy key encountered") else: @@ -1336,13 +1324,9 @@ def sample_grouped_strategies( strategy_impl_cls = _StrategyImpls.MinPSampleOnly case "greedy": strategy_impl_cls = _StrategyImpls.GreedySampleOnly - case ("beam_search", beam_width_in_key, _, early_stopping_key): + case ("beam_search", beam_width_in_key, _, _): beam_width_in = beam_width_in_key - strategy_impl_cls = ( - _StrategyImpls.RegularBeamSearchStep - if early_stopping_key == BeamSearchEarlyStop.TRUE - else _StrategyImpls.CBABeamSearchStep - ) + strategy_impl_cls = _StrategyImpls.CBABeamSearchStep case _: raise NotImplementedError("Unsupported strategy key encountered") if group_logit_indices is None: diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index b3513536f04f..e87080d73797 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -292,31 +292,6 @@ def validate_output(output: GenerationResult, input_prompt: list[int], len(input_prompt), beam_idx) -def validate_disagg_output(output: GenerationResult, input_prompt: list[int], - sampling_params: SamplingParams) -> None: - """Validate disagg beam output without requiring full-run cache history.""" - check_context_logits(output, sampling_params) - assert len(output.outputs) == sampling_params.n - expected_outputs = get_expected_outputs( - input_prompt[-1], num_iterations=sampling_params.max_tokens) - - for beam_idx, beam_output in enumerate(output.outputs): - assert beam_output.finish_reason == "length" - check_generation_logits(beam_output, sampling_params, valid_tokens=None) - check_logprobs(beam_output, sampling_params, valid_tokens=None) - expected_token_ids = expected_outputs.outputs[beam_idx].tolist() - assert beam_output.token_ids == expected_token_ids, ( - f"expected {expected_token_ids} token ids, " - f"got {beam_output.token_ids}") - - assert beam_output.additional_generation_outputs is not None - cache_indirection = beam_output.additional_generation_outputs[ - "cache_indirection"] - assert cache_indirection is not None - assert cache_indirection.shape[1] == sampling_params.best_of - assert cache_indirection.shape[0] == sampling_params.max_tokens - 1 - - def validate_outputs(llm: LLM, input_prompts: list[list[int]], sampling_params: SamplingParams, monkeypatch: pytest.MonkeyPatch, @@ -382,56 +357,6 @@ def _update_requests_hook(self, state: SampleStateTorch, *args, validate_output(output, input_prompts[output_idx], sampling_params) -def validate_disagg_outputs(ctx_llm: LLM, - gen_llm: LLM, - input_prompts: list[list[int]], - sampling_params: SamplingParams, - request_id_base: int = 0) -> None: - """Run context-only then generation-only beam search and validate outputs.""" - - ctx_disagg_params = [ - DisaggregatedParams(request_type="context_only", - disagg_request_id=1000 + request_id_base + idx) - for idx in range(len(input_prompts)) - ] - ctx_outputs = ctx_llm.generate( - deepcopy(input_prompts), - sampling_params=deepcopy(sampling_params), - disaggregated_params=ctx_disagg_params, - use_tqdm=False, - ) - assert isinstance(ctx_outputs, list) - assert len(ctx_outputs) == len(input_prompts) - - gen_disagg_params = [] - for ctx_output in ctx_outputs: - disagg_params = ctx_output.disaggregated_params - assert disagg_params is not None - assert disagg_params.first_gen_tokens is not None - assert disagg_params.first_gen_log_probs is not None - assert len(disagg_params.first_gen_tokens) == sampling_params.best_of - assert len(disagg_params.first_gen_log_probs) == sampling_params.best_of - disagg_params.request_type = "generation_only" - gen_disagg_params.append(disagg_params) - - gen_outputs = gen_llm.generate( - deepcopy(input_prompts), - sampling_params=deepcopy(sampling_params), - disaggregated_params=gen_disagg_params, - use_tqdm=False, - ) - assert isinstance(gen_outputs, list) - assert len(gen_outputs) == len(input_prompts) - for output_idx, output in enumerate(gen_outputs): - validate_disagg_output(output, input_prompts[output_idx], - sampling_params) - - -########################################################################### -# End to end tests -########################################################################### - - @pytest.mark.parametrize("return_log_probs", [True, False]) @pytest.mark.parametrize("gather_generation_logits", [True, False]) @pytest.mark.parametrize("gather_context_logits", [True, False]) @@ -497,10 +422,18 @@ def test_beam_search_disagg_e2e( input_prompts, model_kwargs: dict[str, Any], ) -> None: - if model_kwargs["sampler_type"] != "TorchSampler": - pytest.skip( - "Disaggregated beam score handoff is implemented for TorchSampler") - + """Beam search must be rejected under disaggregated serving. + + Every early_stopping mode keeps a pool of finished candidates, which the + context server can already populate on its first step. That pool 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 holds for both + samplers: TorchSampler serves every mode from the candidate-beams-array + path, and the C++ decoder behind TRTLLMSampler maintains the pool + unconditionally. Rejected at admission until the handoff carries the pool + (TRTLLM-14792). + """ sampling_params = SamplingParams( max_tokens=fixed_params["max_tokens"], n=fixed_params["max_beam_width"], @@ -508,7 +441,6 @@ def test_beam_search_disagg_e2e( use_beam_search=True, end_id=-1, include_stop_str_in_output=True, - additional_model_outputs=["cache_indirection"], ) disagg_kwargs = deepcopy(model_kwargs) @@ -527,38 +459,22 @@ def test_beam_search_disagg_e2e( ), ) - partial_reuse_prompts = [[1, 2, 3], [1, 5, 6]] - - ctx_llm = _build_llm(fixed_params, partial_reuse_prompts, disagg_kwargs) - gen_llm = _build_llm(fixed_params, partial_reuse_prompts, disagg_kwargs) + prompts = [[1, 2, 3]] + ctx_llm = _build_llm(fixed_params, prompts, disagg_kwargs) try: - with ctx_llm, gen_llm: - validate_disagg_outputs(ctx_llm, - gen_llm, - partial_reuse_prompts[:1], - sampling_params, - request_id_base=0) - validate_disagg_outputs(ctx_llm, - gen_llm, - partial_reuse_prompts[1:], - sampling_params, - request_id_base=100) - - # The exhaustive early_stopping modes keep a pool of finished - # candidates that the context server can already populate, but - # that pool is not part of the handoff and is cleared on the - # generation side -- a completion found during the context phase - # would be dropped. The combination is rejected at admission - # until the handoff carries the pool (TRTLLM-14792). - for early_stopping in (0, 2): - exhaustive_params = deepcopy(sampling_params) - exhaustive_params.early_stopping = early_stopping + with ctx_llm: + # Every mode is rejected, including the default (early_stopping + # unset, i.e. True). + for early_stopping in (None, 0, 1, 2): + params = deepcopy(sampling_params) + if early_stopping is not None: + params.early_stopping = early_stopping with pytest.raises(RequestError, match=".*not supported with disaggregated" " serving.*"): _ = ctx_llm.generate( - deepcopy(partial_reuse_prompts[:1]), - sampling_params=exhaustive_params, + deepcopy(prompts), + sampling_params=params, disaggregated_params=[ DisaggregatedParams(request_type="context_only", disagg_request_id=200) @@ -567,7 +483,6 @@ def test_beam_search_disagg_e2e( ) finally: ctx_llm.shutdown() - gen_llm.shutdown() @pytest.mark.parametrize("beam_width", [10]) @@ -2047,6 +1962,40 @@ def test_vbws_cpp_formula_matches_past_array_end(): request, False) == beam_width_array[-1] +@pytest.mark.parametrize( + "beam_width_array, accepted", + [ + ([2, 3, 4], True), + ([4, 4, 4], True), + ([2, 2, 4], True), + ([4, 3, 2], False), + ([2, 4, 3], False), + ], +) +def test_vbws_rejects_decreasing_beam_width_array(beam_width_array: list[int], + accepted: bool): + """Only non-decreasing beam_width_array values are supported. + + The documented VBWS semantics only cover widening (see getBeamWidthByIter + in llmRequest.h), 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 would return beams whose + ancestry and cumulative log-probs are left over from an earlier, wider + step. Reject at admission instead of emitting silently stale beams. + """ + + # Mirrors the check in PyExecutor._validate_request. + def is_accepted(array: list[int]) -> bool: + return not any(b < a for a, b in zip(array, array[1:])) + + assert is_accepted(beam_width_array) == accepted + + request = _vbws_request(beam_width_array) + # The request itself is still constructible; admission is what rejects it, + # and py_beam_width is the array maximum either way. + assert request.py_beam_width == max(beam_width_array) + + def test_vbws_uniform_array_matches_fixed_width(): """A constant beam_width_array must behave exactly like a fixed width.""" max_beam_width = 4 @@ -2374,16 +2323,19 @@ def test_base_class_is_abstract(): 2, 2, 2, torch.ones(1), None, None) @pytest.mark.parametrize( - "early_stopping, expected_cls", + "early_stopping", [ - (BeamSearchEarlyStop.TRUE, _StrategyImpls.RegularBeamSearchStep), - (BeamSearchEarlyStop.FALSE, _StrategyImpls.CBABeamSearchStep), - (BeamSearchEarlyStop.NEVER, _StrategyImpls.CBABeamSearchStep), + BeamSearchEarlyStop.TRUE, + BeamSearchEarlyStop.FALSE, + BeamSearchEarlyStop.NEVER, ], ) @staticmethod @_kernel_test - def test_from_strategies_builds_concrete_impl(early_stopping, expected_cls): + def test_from_strategies_builds_concrete_impl(early_stopping): + # Every stopping mode runs on the candidate-beams-array step; the mode + # only selects the done verdict computed inside it. + expected_cls = _StrategyImpls.CBABeamSearchStep strategies = [_beam_strategy(early_stopping=early_stopping)] impl = expected_cls.from_strategies(strategies, cuda_device=torch.device("cuda")) From cf9325a4d958e2da4e07d60a295c380f684fa09f Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 5 Aug 2026 08:23:48 -0700 Subject: [PATCH 11/36] [TRTLLM-13234][chore] Update comments left stale by the CBA unification 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 --- tensorrt_llm/_torch/pyexecutor/sampler/sampler.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 73db892f2499..2c5341d518e0 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -1989,10 +1989,9 @@ def validate_request(self, request: LlmRequest) -> None: raise ValueError( "Beam search only supports returning the sampled logprob per token" ) - # early_stopping == TRUE (default) is served by the frozen-slot path - # (stopping once all beam slots hold finished beams is equivalent); - # the exhaustive modes by the candidate-beams-array path (see - # beam_search_sampling_batch_cba). + # Every early_stopping mode is served by the candidate-beams-array + # path (see beam_search_sampling_batch_cba); the mode only selects + # the done verdict computed there. @override @nvtx_range("setup_sampler_step") @@ -2073,8 +2072,8 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: if self._use_beam_search: beam_search_store = self.store.beam_search_store assert beam_search_store is not None - # Allocate the CBA tensors on the first exhaustive-early_stopping - # request; beam search with the default mode never needs them. + # Allocate the CBA tensors on the first beam-search request: every + # early_stopping mode runs on that path. if any(_request_uses_cba(request) for request in new_requests): beam_search_store.ensure_cba() _prepare_beam_search( From 2f4727e3d48cde973a9539ca2b2407f86d4546ef Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 5 Aug 2026 18:48:23 -0700 Subject: [PATCH 12/36] [TRTLLM-13234][fix] Key beam-search strategy groups regardless of tuple 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 --- .../_torch/pyexecutor/sampler/beam_search.py | 24 +- .../pyexecutor/sampler/sampler_strategy.py | 12 +- .../_torch/sampler/test_beam_search.py | 297 +++++++++--------- 3 files changed, 168 insertions(+), 165 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py index f4b673660872..9a3eb2cfb2ae 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py @@ -845,19 +845,19 @@ def _take(src: torch.Tensor) -> torch.Tensor: # `min_kept > neg_inf` means the pool holds `caps` finished hypotheses, # i.e. the C++ `numBeamsCBA[slot] >= nBM` test. pool_full = min_kept > neg_inf - if early_stopping == BeamSearchEarlyStop.TRUE: - # HF `True`: stop as soon as `beam_width` finished candidates exist, - # without weighing what is still attainable. Matches C++ - # beamStage3Kernel, which short-circuits to done here. - done = pool_full + if early_stopping != BeamSearchEarlyStop.FALSE: + max_gen = (max_seq_len - prompts.view(-1)).to(cand_len.dtype) + bound_len = torch.where(exponent.view(-1) > 0, max_gen, cand_len.view(-1)) else: - if early_stopping != BeamSearchEarlyStop.FALSE: - max_gen = (max_seq_len - prompts.view(-1)).to(cand_len.dtype) - bound_len = torch.where(exponent.view(-1) > 0, max_gen, cand_len.view(-1)) - else: - bound_len = cand_len.view(-1) - best_attainable = cand_cum[:, 0] / bound_len.to(cand_cum.dtype).pow(exponent.view(-1)) - done = pool_full & (min_kept >= best_attainable) + bound_len = cand_len.view(-1) + best_attainable = cand_cum[:, 0] / bound_len.to(cand_cum.dtype).pow(exponent.view(-1)) + # HF `True` stops as soon as `beam_width` finished candidates exist, without + # weighing what is still attainable -- C++ beamStage3Kernel short-circuits + # to done there. Fold that into the same expression rather than branching on + # it: `early_stopping` is a plain int here, so an extra Python branch is an + # extra Dynamo guard and recompilation of this fullgraph function. + ignore_attainable = early_stopping == BeamSearchEarlyStop.TRUE + done = pool_full & (ignore_attainable | (min_kept >= best_attainable)) # Reorder the finish handler's rolling stop-word window to follow the # beam swap (matching stays correct across swaps). diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py index 6bc94ad6929f..ec9bdc3059fa 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py @@ -1247,7 +1247,17 @@ def strategy_grouping_key(strategy: Strategy) -> _STRATEGY_KEY_TYPE: | ("greedy", None) ): return cast(_STRATEGY_KEY_TYPE, strategy[0]) - case ("beam_search", beam_width_in, beam_width_out, _, _, _, early_stopping): + # Trailing wildcard: row_stride is appended after early_stopping. + case ( + "beam_search", + beam_width_in, + beam_width_out, + _, + _, + _, + early_stopping, + *_, + ): return cast( _STRATEGY_KEY_TYPE, (strategy[0], beam_width_in, beam_width_out, early_stopping), diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index e87080d73797..19c4cdf23596 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -18,7 +18,6 @@ import pathlib as _pl from contextlib import contextmanager, nullcontext from copy import deepcopy -from dataclasses import dataclass from types import SimpleNamespace from typing import Any, Callable, Generator, cast @@ -33,13 +32,13 @@ from tensorrt_llm import LLM, DisaggregatedParams, SamplingParams, TorchLlmArgs from tensorrt_llm._torch.models.checkpoints import HfCheckpointLoader from tensorrt_llm._torch.pyexecutor.llm_request import (LlmRequest, + LlmRequestState, SamplingConfig) from tensorrt_llm._torch.pyexecutor.sampler import (BeamHistory, SampleStateTorch, TorchSampler) -from tensorrt_llm._torch.pyexecutor.sampler.beam_search import _finalize_beam -from tensorrt_llm._torch.pyexecutor.sampler.logprobs import \ - convert_logprobs_tensor_to_list +from tensorrt_llm._torch.pyexecutor.sampler.beam_search import ( + CBAGroupHost, _finalize_beam, _gather_beam_path, _prepare_beam_history_cba) from tensorrt_llm._torch.pyexecutor.sampler.sampler_strategy import ( BEAM_SEARCH_PAD_TOKEN, BeamSearch, BeamSearchEarlyStop, BeamSearchMetadata, CBAState, _StrategyImpls, beam_search_sampling_batch) @@ -1996,6 +1995,30 @@ def is_accepted(array: list[int]) -> bool: assert request.py_beam_width == max(beam_width_array) +def test_beam_strategy_grouping_key_tolerates_trailing_fields(): + """The grouping key must not depend on the BeamSearch tuple's arity. + + strategy_grouping_key() pattern-matches the strategy tuple. It used to + match exactly the fields known at the time, so appending row_stride made + every beam-search strategy fall through to the "Unsupported strategy" + branch and beam search failed for every request. Unit tests missed it + because they build the ops and step classes directly, bypassing the + grouping layer. Pin that trailing fields are ignored. + """ + from tensorrt_llm._torch.pyexecutor.sampler.sampler_strategy import \ + FlashInferGroupedStrategySampler + + strategy = _beam_strategy(early_stopping=BeamSearchEarlyStop.TRUE) + key = FlashInferGroupedStrategySampler.strategy_grouping_key(strategy) + assert key == ("beam_search", strategy.beam_width_in, + strategy.beam_width_out, BeamSearchEarlyStop.TRUE) + + # row_stride is not part of the key: requests differing only in it still + # group together. + other = strategy._replace(row_stride=strategy.row_stride + 1) + assert FlashInferGroupedStrategySampler.strategy_grouping_key(other) == key + + def test_vbws_uniform_array_matches_fixed_width(): """A constant beam_width_array must behave exactly like a fixed width.""" max_beam_width = 4 @@ -2031,163 +2054,133 @@ def test_vbws_dummy_requests_excluded_from_width_check(): "the mixed-beam-width check in ModelEngine and abort the batch") -def test_create_beam_history(): - """Test TorchSampler._create_beam_history method. +def test_gather_beam_path_follows_cache_indirection(): + """Beam ancestry is reconstructed by following cache_indirection. + + Where C++ walks parent pointers back one step at a time, the Python side + keeps cache_indirection as an already-flattened ancestry table and + reconstructs a beam's tokens with a single gather. This pins that gather, + which both finalization paths share. - This test verifies that beam history is correctly reconstructed by following - the cache_indirection backwards to obtain the correct token sequence. + NB: this used to drive _prepare_beam_history end to end. That is now the + CBA path, which merges the finished-candidate pool into the active beams + and reorders the result by normalized score, so a per-beam expectation + computed from cache_indirection alone no longer describes its output. The + merge and ordering are covered by + test_cba_finalize_merges_pool_and_orders_by_score below. """ - @contextmanager - def _uut_provider( - is_warmup: bool) -> Generator[Callable[[], None], None, None]: - test_params = GeneralTestParams() - request = create_default_request(test_params) - sampler = create_default_sampler(test_params) + test_params = GeneralTestParams() + beam_width = test_params.beam_width + num_generated_tokens = test_params.num_generated_tokens - # Extract parameters from the test parameters - beam_width = test_params.beam_width - prompt_len = test_params.prompt_len - num_generated_tokens = test_params.num_generated_tokens - seq_slot = test_params.seq_slot - vocab_size = test_params.vocab_size - num_logprobs = test_params.num_logprobs + 1 - beam_search_store = sampler.store.beam_search_store - assert beam_search_store is not None - cache_indirection = beam_search_store.cache_indirection - assert cache_indirection is not None - original_tokens = beam_search_store.original_tokens - assert original_tokens is not None - original_logprobs = torch.zeros( - (beam_width, num_generated_tokens, num_logprobs), - dtype=torch.float32, - device=original_tokens.device) - original_logprob_indices = torch.zeros( - (beam_width, num_generated_tokens, num_logprobs), - dtype=torch.int32, - device=original_tokens.device) - original_cum_logprobs = beam_search_store.cum_log_probs - assert original_cum_logprobs is not None + torch.manual_seed(42) + # current_path[beam, t] is the token beam `beam` held at position t before + # correction; cache_indirection[beam, t] names the beam it descended from + # at that position. + current_path = torch.randint(0, + test_params.vocab_size, + (beam_width, num_generated_tokens), + dtype=torch.int32, + device="cuda") + cache_indirection = torch.randint(0, + beam_width, + (beam_width, num_generated_tokens), + dtype=torch.int64, + device="cuda") + + corrected = _gather_beam_path(current_path=current_path, + cache_indirection=cache_indirection) + + expected = torch.zeros_like(current_path) + for beam in range(beam_width): + for t in range(num_generated_tokens): + expected[beam, t] = current_path[cache_indirection[beam, t], t] + torch.testing.assert_close(corrected, expected) + + # An identity table must leave every beam untouched. + identity = (torch.arange(beam_width, device="cuda", dtype=torch.int64).view( + -1, 1).expand(-1, num_generated_tokens).contiguous()) + torch.testing.assert_close( + _gather_beam_path(current_path=current_path, + cache_indirection=identity), current_path) - # Fill the request with some random tokens that will be overwritten by the beam search sampling - # Beam history is created before add_token is called - request.set_generated_tokens( - torch.randint(0, - vocab_size, (beam_width, num_generated_tokens - 1), - dtype=torch.int32).tolist()) - # random fill - torch.manual_seed(42) - original_tokens[seq_slot, :beam_width, prompt_len:prompt_len + - num_generated_tokens] = torch.randint( - 0, - beam_width, (beam_width, num_generated_tokens), - dtype=torch.int32) - assert original_tokens.sum( - ) > 0, "Original tokens must not only contain zeros. Otherwise change the seed." - - original_logprobs[:beam_width] = torch.randn( - (beam_width, num_generated_tokens, original_logprobs.shape[-1]), - dtype=torch.float32) - original_logprob_indices[:beam_width] = torch.randint( - 0, - vocab_size, - (beam_width, num_generated_tokens, original_logprobs.shape[-1]), - dtype=torch.int32) - assert (original_logprobs != 0).sum( - ) > 0, "Original log probs must not only contain zeros. Otherwise change the seed." - assert (original_logprob_indices).sum( - ) > 0, "Original log prob indices must not only contain zeros. Otherwise change the seed." - - # set the logprobs in the request: - token_logprobs = convert_logprobs_tensor_to_list( - original_logprob_indices[:beam_width, :num_generated_tokens - 1], - original_logprobs[:beam_width, :num_generated_tokens - 1], - ) - request.py_result.set_log_probs( - token_logprobs, - cum_log_probs=torch.zeros_like( - original_cum_logprobs[seq_slot, :beam_width]).tolist()) - - original_cum_logprobs[seq_slot, :beam_width] = torch.randn( - (beam_width, ), dtype=torch.float32) - assert (original_cum_logprobs != 0).sum( - ) > 0, "Original cumulative log probs must not only contain zeros. Otherwise change the seed." - - cache_indirection[seq_slot, :beam_width, prompt_len:prompt_len + - num_generated_tokens] = torch.randint( - 0, - beam_width, (beam_width, num_generated_tokens), - dtype=torch.int32) - assert cache_indirection[ - seq_slot, :beam_width, - prompt_len:prompt_len + num_generated_tokens].sum( - ) > 0, "Deterministic offsets must not only contain zeros. Otherwise change the seed." - - # set the new log probs and tokens for the beam search sampling - log_probs_store = sampler.store.log_probs_store - log_probs_store.sampled_log_probs[ - seq_slot, :beam_width] = original_logprobs[:beam_width, - num_generated_tokens - 1, - 0:1] - sampler.store.new_tokens[ - 0, seq_slot, : - beam_width] = original_logprob_indices[:beam_width, - num_generated_tokens - 1, 0] - - @dataclass - class UutResult: - beam_history_builder: Callable[[], BeamHistory | None] | None - - @dataclass - class UutResultWrapper: - result: UutResult | None = None - - res = UutResultWrapper() - # test - def _uut(res=res): - res.result = UutResult( - beam_history_builder=sampler._beam_search._prepare_beam_history( - request, - finish_reasons=torch.ones((beam_width, ), dtype=torch.int), - d2h_copier=sampler._copy_to_host, - ), ) +def test_cba_finalize_merges_pool_and_orders_by_score(): + """CBA finalization ranks pool entries against the live beams. - yield _uut + _prepare_beam_history_cba concatenates the finished-candidate pool with + the active beams, orders the union by normalized score and emits the top + num_beams. Pool entries carry their own (shorter) length and are padded to + the output width; active beams contribute the full generated window. This + pins that selection, which the plain ancestry gather in + test_gather_beam_path_follows_cache_indirection does not cover. + """ + num_beams = 3 + num_generated = 4 + pad = BEAM_SEARCH_PAD_TOKEN + + request = _vbws_request(None, max_beam_width=num_beams) + request.state = LlmRequestState.GENERATION_IN_PROGRESS + request.py_decoding_iter = num_generated + request.decoding_iter = num_generated + request.py_seq_slot = 0 + prompt_len = request.py_prompt_len + # num_generated_tokens is derived from the request's token count, so give + # it the generated tokens the beam state below describes. The last token + # of the step is not added yet, hence num_generated - 1. + request.set_generated_tokens([[0] * (num_generated - 1)] * num_beams) + + total = prompt_len + num_generated + # Identity ancestry: each active beam keeps its own tokens, so the ordering + # rather than the gather is what this test observes. + cache_indirection = (torch.arange(num_beams, + dtype=torch.int64).view(-1, 1).expand( + -1, total).contiguous().unsqueeze(0)) + original_tokens = torch.zeros((1, num_beams, total), dtype=torch.int32) + original_tokens[0, :, prompt_len:] = torch.tensor( + [[31, 32, 33, 34], [41, 42, 43, 44], [51, 52, 53, 54]], + dtype=torch.int32) - torch.cuda.synchronize() - assert res.result is not None - beam_history_builder = res.result.beam_history_builder - assert beam_history_builder is not None - beam_history = beam_history_builder() - assert beam_history is not None - - # expected selection: - # Currently beam history only contains the generated tokens, not the prompt tokens. - expected_tokens = torch.zeros( - (sampler.max_beam_width, num_generated_tokens), dtype=torch.int32) - expected_logprobs = torch.zeros( - (beam_width, num_generated_tokens, original_logprobs.shape[-1]), - dtype=torch.float32) - for gen_idx in range(num_generated_tokens): - token_idx = prompt_len + gen_idx - expected_tokens[:, gen_idx] = original_tokens[ - seq_slot, cache_indirection[seq_slot, :, token_idx], token_idx] - expected_logprobs[:, gen_idx] = original_logprobs[cache_indirection[ - seq_slot, :beam_width, token_idx], gen_idx] - - torch.testing.assert_close(beam_history.tokens[:beam_width], - expected_tokens[:beam_width]) - # test logprobs as well - assert beam_history.logprobs is not None - torch.testing.assert_close(beam_history.logprobs[:beam_width], - expected_logprobs[:beam_width]) - assert beam_history.cum_logprobs is not None - torch.testing.assert_close( - beam_history.cum_logprobs[:beam_width], - original_cum_logprobs[seq_slot, :beam_width].to("cpu")) + # Pool: entry 0 outranks every live beam, entry 1 sits between them, + # entry 2 is an unused (-inf) slot that must never be selected. + cba_tokens = torch.full((1, num_beams, total), pad, dtype=torch.int32) + cba_tokens[0, 0, :2] = torch.tensor([11, 12], dtype=torch.int32) + cba_tokens[0, 1, :3] = torch.tensor([21, 22, 23], dtype=torch.int32) + + cba_group = CBAGroupHost( + pos={0: 0}, + should_stop=torch.tensor([True]), + cache_indirection=cache_indirection, + original_tokens=original_tokens, + cum=torch.tensor([[7.0, 5.0, 3.0]]), + cba_tokens=cba_tokens, + cba_cum=torch.tensor([[90.0, 6.0, 0.0]]), + cba_normed=torch.tensor([[90.0, 6.0, float("-inf")]]), + cba_lengths=torch.tensor([[2, 3, 0]], dtype=torch.int32), + original_log_probs=None, + cba_log_probs=None, + ) - run_test_with_warmup(_uut_provider, max_sync_s=1) + builder = _prepare_beam_history_cba(request, cba_group=cba_group) + assert builder is not None + history = builder() + assert history is not None + + # Ranking over the union: pool0=90 > active0=7 > pool1=6, so the second + # pool entry and the two weaker live beams drop out. + torch.testing.assert_close( + history.tokens, + torch.tensor( + [ + [11, 12, pad, pad], # pool entry, padded past its length + [31, 32, 33, 34], # best live beam, full window + [21, 22, 23, pad], # next pool entry + ], + dtype=torch.int32)) + assert history.cum_logprobs is not None + torch.testing.assert_close(history.cum_logprobs, + torch.tensor([90.0, 7.0, 6.0])) def test_finish_beams(): From d83e91a977368d8030d3ff05cf877e7ee1e47807 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 5 Aug 2026 19:05:37 -0700 Subject: [PATCH 13/36] [TRTLLM-13234][fix] Stop the beam padding sentinel reaching request histories 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 --- .../_torch/pyexecutor/model_engine.py | 12 +++++-- .../_torch/pyexecutor/sampler/sampler.py | 10 +++++- .../pyexecutor/sampler/sampler_strategy.py | 14 +++++++-- .../_torch/sampler/test_beam_search.py | 31 +++++++++++++++++++ 4 files changed, 62 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index 865bed3b5d4b..3715962c18a4 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -7873,9 +7873,17 @@ def _execute_logit_post_processors(self, for request in requests: if is_context_request: beam_width = 1 + row_stride = 1 else: + # Generation rows are laid out at the static admission + # width, so that is the stride between requests, while + # only the leading beam_width rows hold live beams under + # a variable beam width array. Advancing the offset by the + # narrower width would make every request after the first + # rewrite another request's logits rows in place. beam_width = request.get_beam_width_by_iter( for_next_iteration=False) + row_stride = request.py_beam_width logits_processors = getattr(request, "py_logits_post_processors", None) @@ -7888,13 +7896,13 @@ def _execute_logit_post_processors(self, if (is_context_request and request.py_orig_prompt_len < len(token_ids[0])): # Skip as we only need to apply logit processor on the last context request - logits_row_offset += beam_width + logits_row_offset += row_stride continue self._apply_logits_processors(request, logits_processors, logits_tensor, beam_width, token_ids, logits_row_offset) - logits_row_offset += beam_width + logits_row_offset += row_stride def wait_for_input_copy(self): """ diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 2c5341d518e0..deddd9001079 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -106,6 +106,7 @@ DEFAULT_BEAM_IDX, DEFAULT_STEP_IDX, FinishReasonsList, + _get_beam_width_out, _get_max_beam_width, _request_get_sampling_params, add_token, @@ -2436,7 +2437,14 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None: if (beam_history := _maybe_build_beam_history(req_idx)) is not None: _finalize_beam(req, beam_history) else: - for beam_idx in range(req.py_beam_width): + # Only the leading beam_width_out columns hold real tokens; + # the op pads the rest of the store-width row with + # BEAM_SEARCH_PAD_TOKEN. Appending those would put the + # sentinel into the request's token history, which is + # visible to streaming consumers and to anything reading + # get_tokens() mid-flight (e.g. the token-ban suffix + # matching) even though finalization later rewrites it. + for beam_idx in range(_get_beam_width_out(req)): # Beam search does not support speculative decoding. add_token(req, new_tokens_list, beam_idx=beam_idx) self.handle_logprobs(req, logprobs_state_list=logprobs_state_list, count=1) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py index ec9bdc3059fa..14255d25e4cc 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py @@ -1127,7 +1127,11 @@ def sample( seeds: Optional["RequestSeeds"] = None, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: assert group_metadata is not None and isinstance(group_metadata, BeamSearchMetadata) - temperature = self._temperature.repeat_interleave(self._beam_width_in) + # Temperature is applied before the op slices the padding rows off, + # so it must cover every row the forward path laid out: the static + # admission width, which exceeds beam_width_in while a variable + # beam width array is still widening. + temperature = self._temperature.repeat_interleave(self._row_stride) logits = self._prepare_logits_with_temperature(logits, group_logit_indices, temperature) return self._select_and_update(logits, group_metadata) @@ -1340,7 +1344,13 @@ def sample_grouped_strategies( case _: raise NotImplementedError("Unsupported strategy key encountered") if group_logit_indices is None: - assert logits.size(0) == beam_width_in * len(strategies) + # Beam-search rows are laid out at the static admission width + # (row_stride), which exceeds beam_width_in on a widening + # variable-beam-width step; the op slices down to the live beams. + rows_per_request = beam_width_in + if strategies and strategies[0][0] == "beam_search": + rows_per_request = cast(BeamSearch, strategies[0]).row_stride + assert logits.size(0) == rows_per_request * len(strategies) else: assert group_logit_indices.size(0) == beam_width_in * len(strategies) strategy_impl = strategy_impl_cls.from_strategies(strategies, cuda_device=logits.device) diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 19c4cdf23596..06c3db305cbb 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -613,6 +613,31 @@ def recording_get_beam_width_by_iter(self: LlmRequest, monkeypatch.setattr(LlmRequest, "get_beam_width_by_iter", recording_get_beam_width_by_iter) + # Inspect the requests right after every update_requests(): the sampling + # op pads its output row out to the store's beam width, and appending + # those columns would put BEAM_SEARCH_PAD_TOKEN into the request's token + # history. Checking only the final outputs cannot see it, because + # finalization rewrites every beam from the corrected paths -- but the + # padded history is visible to streaming consumers and to anything reading + # get_tokens() mid-flight. + padded_histories: list[tuple[int, int, list[int]]] = [] + unwrapped_update_requests = TorchSampler.update_requests + + def recording_update_requests(self: TorchSampler, state, *args, **kwargs): + result = unwrapped_update_requests(self, state, *args, **kwargs) + for req in state.requests: + if req.py_beam_width <= 1: + continue + for beam_idx in range(req.py_beam_width): + tokens = list(req.get_tokens(beam_idx)) + if BEAM_SEARCH_PAD_TOKEN in tokens: + padded_histories.append( + (req.py_decoding_iter, beam_idx, tokens)) + return result + + monkeypatch.setattr(TorchSampler, "update_requests", + recording_update_requests) + gc.collect(2) # force destruction of any other LLM instances with _single_process_context(): llm = LLM( @@ -672,6 +697,12 @@ def recording_get_beam_width_by_iter(self: LlmRequest, f"expected every width in {beam_width_array} to be exercised, " f"but only saw {sorted(set(actual))}") + # No step may leave the padding sentinel in a request's token history. + assert not padded_histories, ( + "BEAM_SEARCH_PAD_TOKEN leaked into the request token history; " + "update_requests() must append only the beams the step produced, not " + f"the full store width. First offenders: {padded_histories[:3]}") + ########################################################################### # Unit tests From 3e162e5f22a86c781ca7acc6b895bea62db04656 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 5 Aug 2026 19:39:30 -0700 Subject: [PATCH 14/36] [TRTLLM-13234][chore] Refresh comments the CBA unification invalidated 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 --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 10 +- .../_torch/pyexecutor/sampler/beam_search.py | 53 ++-- .../_torch/pyexecutor/sampler/sampler.py | 2 +- .../pyexecutor/sampler/sampler_strategy.py | 30 +-- .../_torch/sampler/test_beam_search.py | 236 +----------------- 5 files changed, 58 insertions(+), 273 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 2fd71afded98..cefea423f0bf 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -6412,7 +6412,15 @@ def _prepare_disagg_gen_transmission_complete(self, scheduled_batch): def _update_sampler_state_for_disagg_gen_request(self, req, beam_width, first_gen_tokens) -> bool: - """Update beam sampler state with context-side first-token data.""" + """Update beam sampler state with context-side first-token data. + + NB: currently unreachable past the beam_width <= 1 guard. Admission + rejects beam search under disaggregated serving (see + _validate_request), because the finished-candidate pool the context + server builds is not part of the handoff. Kept rather than deleted: + this is the seeding half of that handoff and becomes live again once + the pool is transferred; TRTLLM-14792. + """ if beam_width <= 1: return True diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py index 9a3eb2cfb2ae..5a1b3676ea47 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py @@ -108,8 +108,7 @@ class _CBAFields: same tensors). Field shapes and semantics are documented here once; both subclasses inherit these fields. - Only maintained for requests using exhaustive early_stopping modes - (early_stopping != TRUE); allocated lazily by + Maintained for every beam-search request; allocated lazily by ``BeamSearchStore.ensure_cba`` on the first such request. Fields written on every beam-search step regardless of stopping mode @@ -144,7 +143,7 @@ class _CBAFields: @dataclass(kw_only=True) class CBAState(_CBAFields): """Candidate-Beams-Array (CBA) state, present only for requests using - exhaustive early_stopping modes (early_stopping != TRUE); see + every beam-search request, whatever its early_stopping mode; see beam_search_sampling_batch_cba. A per-step view over the persistent ``BeamSearchStore`` (shared CBA tensors @@ -200,7 +199,7 @@ class BeamSearchStore: for stop word detection.""" seq_offsets: torch.Tensor """[max_num_sequences] int64, cached ``arange(max_num_sequences) * - max_beam_width`` used by ``beam_search_sampling_batch`` to flatten + max_beam_width`` used by ``beam_search_sampling_batch_cba`` to flatten (batch_idx, beam_idx) pairs.""" beam_idx_arange: torch.Tensor """[max_beam_width] int32, cached ``arange(max_beam_width)`` used as the @@ -220,7 +219,7 @@ class BeamSearchStore: batch_dones: torch.Tensor """[max_num_sequences] bool, per-slot beam-search termination verdict.""" cba: Optional[_CBAFields] = None - """CBA tensors; None until the first exhaustive-early_stopping request.""" + """CBA tensors; None until the first beam-search request.""" @classmethod def create( @@ -255,8 +254,8 @@ def create( def ensure_cba(self) -> _CBAFields: """Allocate the CBA tensors on first use and return them. - Called when a request with an exhaustive early_stopping mode is - admitted; idempotent afterwards. + Called when the first beam-search request is admitted, whatever its + early_stopping mode; idempotent afterwards. """ if self.cba is None: shape = self.original_tokens.shape @@ -325,7 +324,7 @@ def _gather_beam_path( @dataclass(kw_only=True) class BeamSearchMetadata(StrategyMetadata): - """Stateful tensors required by beam_search_sampling_batch.""" + """Stateful tensors required by beam_search_sampling_batch_cba.""" cache_indirection: torch.Tensor cache_indirection_buffer: torch.Tensor @@ -347,8 +346,8 @@ class BeamSearchMetadata(StrategyMetadata): skipped — multi-token stop-word matching would then be unreliable across beam swaps.""" cba: Optional[CBAState] = None - """Candidate-Beams-Array state, present only for exhaustive early_stopping - modes (early_stopping != TRUE); None for the regular beam-search path.""" + """Candidate-Beams-Array state, present for every beam-search request + regardless of its early_stopping mode; None until the first one.""" def _update_cache_indirection_buffer( @@ -413,9 +412,15 @@ def _pad_next_tokens(next_tokens: torch.Tensor, store_width: int) -> torch.Tenso """Pad a [batch, beam_width_out] token tensor to the store's beam width. The batched sampling buffers are allocated at the maximum beam width; on - variable-beam-width steps the op produces fewer columns, and the padding - (BEAM_SEARCH_PAD_TOKEN) is never consumed — finalization rewrites all - beam tokens from the corrected paths. + variable-beam-width steps the op produces fewer columns and the rest are + filled with BEAM_SEARCH_PAD_TOKEN. + + Consumers must read only the leading ``beam_width_out`` columns. Appending + the padded ones puts the sentinel into the request's token history, which + finalization does later overwrite from the corrected paths -- but the + padded history is visible in the meantime to streaming consumers and to + anything reading ``get_tokens()`` mid-flight. See the + ``_get_beam_width_out`` bound in ``TorchSampler.update_requests``. """ if next_tokens.size(1) >= store_width: return next_tokens @@ -534,6 +539,12 @@ def beam_search_sampling_batch( ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: """Sample beam_width tokens for each request in parallel. + NB: no longer dispatched to. Every early_stopping mode now runs on + ``beam_search_sampling_batch_cba``; this pool-free variant freezes a + finished beam in its slot instead of vacating it, which diverges from the + C++ decoder and from vLLM. Kept as the restore hook referenced by + ``_request_uses_cba``, and still exercised directly by unit tests. + ``length_penalty`` normalizes the beam-selection ranking key as ``cum_log_prob / gen_length**length_penalty``, and ``diversity_rate`` adds ``diversity_rate * source_beam_index`` to it; see ``beam_candidate_topk``. The stored ``cum_log_probs`` remain @@ -902,8 +913,8 @@ def beam_search_sampling_batch_cba( diversity_rate: "torch.Tensor | float | None" = None, return_probs: bool = True, ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - """Beam-search step with a candidate-beams array (CBA) for the exhaustive - early_stopping modes (``early_stopping != TRUE``): + """Beam-search step with a candidate-beams array (CBA). Every + early_stopping mode runs here; the mode only selects the done verdict: - The top ``2 * beam_width`` expansion candidates (ranked by raw cumulative log-prob, plus the optional diversity adjustment — the @@ -933,7 +944,7 @@ def beam_search_sampling_batch_cba( """ args = beam_search_args cba = args.cba - assert cba is not None, "CBA metadata is required for early_stopping != TRUE" + assert cba is not None, "CBA metadata is required for beam search" num_beams = beam_width_out device = logits.device slots = args.seq_slots @@ -1161,7 +1172,7 @@ def _prepare_beam_search( beam_search_store.beam_gen_lengths.index_fill_(0, seq_slots_long, 0) beam_search_store.prompt_lens.index_copy_(0, seq_slots_long, prompt_lens_cuda) beam_search_store.batch_dones.index_fill_(0, seq_slots_long, False) - # The CBA tensors only exist once an exhaustive-early_stopping request has + # The CBA tensors only exist once a beam-search request has # been admitted; nothing reads them before that, so skip the reset. cba = beam_search_store.cba if cba is not None: @@ -1402,6 +1413,10 @@ def _check_beam_search_stop_criteria( ) -> torch.Tensor: """Check if the stop criteria is met for the request. + NB: only reachable from the pool-free finalize branch, which nothing + dispatches to any more -- the CBA path computes its own verdict in + ``prepare_cba_group_host``. Kept alongside ``beam_search_sampling_batch``. + Returns a boolean tensor of shape (), whose value is computed asynchronously. """ return (finish_reasons[: request.py_beam_width] > 0).sum() == request.py_beam_width @@ -1588,6 +1603,10 @@ def _prepare_beam_history( assert cba_group is not None return _prepare_beam_history_cba(request, cba_group=cba_group) + # Everything below is the pool-free finalize, unreachable while + # _request_uses_cba is unconditionally true. Kept with + # beam_search_sampling_batch as the restore path. + # Gather data used for skipping beam history processing need_finalize_due_to_stop_words = self._has_multi_token_stop_words(request) if need_finalize_due_to_stop_words: diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index deddd9001079..e1cbc7725577 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -2269,7 +2269,7 @@ def _add_metadata_to_grouped_requests( beam_idx_arange=beam_search_store.beam_idx_arange, beam_gen_lengths=beam_search_store.beam_gen_lengths, stop_past_tokens=self._finish_reasons_handler.store.past_tokens_cuda, - # None unless an exhaustive-early_stopping request has been + # None unless a beam-search request has been # admitted; the CBA tensors are not allocated before that. cba=None if beam_search_store.cba is None diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py index 14255d25e4cc..848d6ed90164 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py @@ -1014,8 +1014,8 @@ class BeamSearchStep(StrategyImpl): ``sample`` applies the shared temperature preprocessing and delegates to the ``_select_and_update`` hook, implemented per stopping mode by - ``RegularBeamSearchStep`` (early_stopping == TRUE) and - ``CBABeamSearchStep`` (FALSE / NEVER). With-probs is a constructor flag + ``CBABeamSearchStep``, which every stopping mode uses -- the mode only + selects the done verdict computed inside it. With-probs is a constructor flag (``computes_probs``), not a subclass. """ @@ -1141,27 +1141,13 @@ def _select_and_update( ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: """Mode-specific candidate selection and state update.""" - class RegularBeamSearchStep(BeamSearchStep): - """early_stopping == TRUE: the regular beam-search step.""" - - @override - def _select_and_update( - self, logits: torch.Tensor, group_metadata: BeamSearchMetadata - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - return beam_search_sampling_batch( - logits, - beam_width_in=self._beam_width_in, - beam_width_out=self._beam_width_out, - row_stride=self._row_stride, - beam_search_args=group_metadata, - temperature=None, - length_penalty=self._length_penalty, - diversity_rate=self._diversity_rate, - return_probs=self.computes_probs(), - ) - class CBABeamSearchStep(BeamSearchStep): - """early_stopping in {FALSE, NEVER}: the candidate-beams-array step.""" + """The candidate-beams-array step, used by every early_stopping mode. + + The mode only selects the done verdict computed inside the op: TRUE + stops as soon as the pool is full, FALSE and NEVER additionally weigh + what is still attainable. + """ def __init__( self, diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 06c3db305cbb..fc8765f4d4ba 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -18,7 +18,6 @@ import pathlib as _pl from contextlib import contextmanager, nullcontext from copy import deepcopy -from types import SimpleNamespace from typing import Any, Callable, Generator, cast import pytest @@ -46,8 +45,7 @@ from tensorrt_llm.bindings.internal.batch_manager import \ LlmRequest as CppLlmRequest from tensorrt_llm.executor import RequestError -from tensorrt_llm.executor.result import (CompletionOutput, GenerationResult, - Logprob) +from tensorrt_llm.executor.result import CompletionOutput, GenerationResult from tensorrt_llm.llmapi import (CacheTransceiverConfig, CudaGraphConfig, KvCacheConfig) @@ -1623,232 +1621,6 @@ def test_beam_search_cba_reorders_stop_window(): assert stop_window[:, 0, 1].tolist() == [200, 200, 200] -@_kernel_test -def test_beam_search_sampling_batch_reorders_stop_window(): - """The ES=1 path must also reorder the stop-word window on beam swaps.""" - batch_size = 1 - beam_width = 2 - vocab_size = 6 - max_batch_size = 1 - seq_len = 6 - - seq_slots = torch.arange(batch_size, dtype=torch.int64) - stop_window = torch.zeros((3, max_batch_size, beam_width), - dtype=torch.int32) - stop_window[:, 0, 0] = 100 - stop_window[:, 0, 1] = 200 - metadata = BeamSearchMetadata( - cache_indirection=torch.zeros((max_batch_size, beam_width, seq_len + 1), - dtype=torch.int32), - cache_indirection_buffer=torch.full( - (max_batch_size, beam_width, seq_len + 1), -1, dtype=torch.int32), - cum_log_probs=torch.zeros((max_batch_size, beam_width), - dtype=torch.float32), - seq_slots=seq_slots, - seq_lens=torch.full((batch_size, ), seq_len, dtype=torch.int32), - finished_beams=torch.zeros((max_batch_size, beam_width), - dtype=torch.int32), - new_log_probs=torch.zeros((max_batch_size, beam_width), - dtype=torch.float32), - predecessor_beams=torch.zeros((max_batch_size, beam_width), - dtype=torch.int32), - seq_offsets=torch.arange(max_batch_size, dtype=torch.int64) * - beam_width, - beam_idx_arange=torch.arange(beam_width, dtype=torch.int32), - beam_gen_lengths=torch.zeros((max_batch_size, beam_width), - dtype=torch.int32), - stop_past_tokens=stop_window, - ) - metadata.cum_log_probs[0] = torch.tensor([-5.0, -1.0]) - - # beam 1 dominates: both slots descend from beam 1 - logits = torch.full((batch_size * beam_width, vocab_size), -50.0) - logits[1, 2] = 10.0 - logits[1, 3] = 9.0 - - beam_search_sampling_batch( - logits=logits, - beam_width_in=beam_width, - beam_width_out=beam_width, - beam_search_args=metadata, - temperature=1.0, - return_probs=False, - ) - assert metadata.predecessor_beams[0].tolist() == [1, 1] - assert stop_window[:, 0, 0].tolist() == [200, 200, 200] - assert stop_window[:, 0, 1].tolist() == [200, 200, 200] - - -@_kernel_test -def test_beam_search_sampling_batch_disagg_handoff(): - """Test context-first disagg beam handoff seeds gen-side beam scores.""" - - test_params = GeneralTestParams() - batch_size = test_params.batch_size - beam_width = test_params.beam_width - vocab_size = test_params.vocab_size - max_batch_size = test_params.max_batch_size - prompt_len = test_params.prompt_len - temperature = 1.0 - - seq_slots = torch.arange( - batch_size, dtype=torch.int64) + (max_batch_size - batch_size) // 2 - seq_offsets = torch.arange(max_batch_size, dtype=torch.int64) * beam_width - beam_idx_arange = torch.arange(beam_width, dtype=torch.int32) - - def make_metadata(seq_len: int) -> BeamSearchMetadata: - return BeamSearchMetadata( - cache_indirection=torch.zeros( - (max_batch_size, beam_width, prompt_len + 2), - dtype=torch.int32), - cache_indirection_buffer=torch.full( - (max_batch_size, beam_width, prompt_len + 2), - -1, - dtype=torch.int32), - cum_log_probs=torch.zeros((max_batch_size, beam_width), - dtype=torch.float32), - seq_slots=seq_slots, - seq_lens=torch.full((batch_size, ), seq_len, dtype=torch.int32), - finished_beams=torch.zeros((max_batch_size, beam_width), - dtype=torch.int32), - new_log_probs=torch.zeros((max_batch_size, beam_width), - dtype=torch.float32), - predecessor_beams=torch.zeros((max_batch_size, beam_width), - dtype=torch.int32), - seq_offsets=seq_offsets, - beam_idx_arange=beam_idx_arange, - beam_gen_lengths=torch.zeros((max_batch_size, beam_width), - dtype=torch.int32), - ) - - torch.manual_seed(43) - context_logits = torch.randn((batch_size, vocab_size), dtype=torch.float32) - generation_logits = torch.randn((batch_size * beam_width, vocab_size), - dtype=torch.float32) - for req_idx in range(batch_size): - context_logits[req_idx, req_idx * beam_width:req_idx * beam_width + - beam_width] += torch.tensor([8.0, 7.0, 1.0]) - for beam_idx in range(beam_width): - row = req_idx * beam_width + beam_idx - generation_logits[row, 10 + beam_idx] += 4.0 + beam_idx - - # Baseline: one continuous request. Context produces first beams, then - # generation continues with beam_width input beams. - continuous_metadata = make_metadata(prompt_len) - first_tokens, _ = beam_search_sampling_batch( - logits=context_logits, - beam_width_in=1, - beam_width_out=beam_width, - beam_search_args=continuous_metadata, - temperature=temperature, - return_probs=False, - ) - first_gen_scores = continuous_metadata.cum_log_probs[ - seq_slots, :beam_width].clone() - assert first_tokens.shape == (batch_size, beam_width) - - continuous_metadata.seq_lens = torch.full((batch_size, ), - prompt_len + 1, - dtype=torch.int32) - continuous_next_tokens, _ = beam_search_sampling_batch( - logits=generation_logits, - beam_width_in=beam_width, - beam_width_out=beam_width, - beam_search_args=continuous_metadata, - temperature=temperature, - return_probs=False, - ) - continuous_cum_log_probs = continuous_metadata.cum_log_probs[ - seq_slots, :beam_width].clone() - continuous_new_log_probs = continuous_metadata.new_log_probs[ - seq_slots, :beam_width].clone() - continuous_predecessor_beams = continuous_metadata.predecessor_beams[ - seq_slots, :beam_width].clone() - continuous_cache_indirection = continuous_metadata.cache_indirection[ - seq_slots, :beam_width, :prompt_len + 2].clone() - - # Disaggregated generation starts from reset sampler buffers, then seeds - # the first-token beam scores derived from context first_gen_log_probs. - from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor - - disagg_metadata = make_metadata(prompt_len + 1) - original_tokens = torch.zeros((max_batch_size, beam_width, prompt_len + 2), - dtype=torch.int32) - fake_executor = cast( - PyExecutor, - SimpleNamespace(sampler=SimpleNamespace(store=SimpleNamespace( - beam_search_store=SimpleNamespace( - original_tokens=original_tokens, - cache_indirection=disagg_metadata.cache_indirection, - cum_log_probs=disagg_metadata.cum_log_probs, - beam_gen_lengths=disagg_metadata.beam_gen_lengths, - ))))) - for req_idx, seq_slot in enumerate(seq_slots.tolist()): - first_gen_tokens = first_tokens[req_idx, :beam_width].tolist() - first_gen_log_probs = [{ - token_id: - Logprob(logprob=float(first_gen_scores[req_idx, beam_idx].item())) - } for beam_idx, token_id in enumerate(first_gen_tokens)] - req = SimpleNamespace( - py_seq_slot=seq_slot, - py_request_id=req_idx, - prompt_len=prompt_len, - py_disaggregated_params=DisaggregatedParams( - request_type="generation_only", - first_gen_log_probs=first_gen_log_probs, - ), - ) - PyExecutor._update_sampler_state_for_disagg_gen_request( - fake_executor, req, beam_width, first_gen_tokens) - - expected_beam_indices = torch.arange(beam_width, dtype=torch.int32).expand( - batch_size, beam_width) - torch.testing.assert_close( - original_tokens[seq_slots, :beam_width, prompt_len], first_tokens) - torch.testing.assert_close( - disagg_metadata.cache_indirection[seq_slots, :beam_width, prompt_len], - expected_beam_indices) - torch.testing.assert_close( - disagg_metadata.cum_log_probs[seq_slots, :beam_width], first_gen_scores) - disagg_next_tokens, _ = beam_search_sampling_batch( - logits=generation_logits, - beam_width_in=beam_width, - beam_width_out=beam_width, - beam_search_args=disagg_metadata, - temperature=temperature, - return_probs=False, - ) - - torch.testing.assert_close(disagg_next_tokens, continuous_next_tokens) - torch.testing.assert_close( - disagg_metadata.cache_indirection[seq_slots, :beam_width, :prompt_len + - 2], continuous_cache_indirection) - torch.testing.assert_close( - disagg_metadata.cum_log_probs[seq_slots, :beam_width], - continuous_cum_log_probs) - torch.testing.assert_close( - disagg_metadata.new_log_probs[seq_slots, :beam_width], - continuous_new_log_probs) - torch.testing.assert_close( - disagg_metadata.predecessor_beams[seq_slots, :beam_width], - continuous_predecessor_beams) - - # This is the regression guard for disagg: if gen-side cum_log_probs remain - # reset after receiving first_gen_tokens, the next step is scored incorrectly. - unseeded_metadata = make_metadata(prompt_len + 1) - beam_search_sampling_batch( - logits=generation_logits, - beam_width_in=beam_width, - beam_width_out=beam_width, - beam_search_args=unseeded_metadata, - temperature=temperature, - return_probs=False, - ) - assert not torch.allclose( - unseeded_metadata.cum_log_probs[seq_slots, :beam_width], - continuous_cum_log_probs) - - def create_default_request(test_params: GeneralTestParams) -> LlmRequest: sampling_params = SamplingParams(n=test_params.beam_width, best_of=test_params.beam_width, @@ -2502,9 +2274,9 @@ def test_exhaustive_early_stopping_allowed_without_disagg( sampler_type: str, early_stopping: int, ): - # The exhaustive early_stopping modes are only rejected for - # disaggregated serving (the finished-candidate pool is not part of the - # handoff; TRTLLM-14792). Regular serving must still accept them. + # Beam search is rejected wholesale under disaggregated serving (the + # finished-candidate pool is not part of the handoff; TRTLLM-14792), + # but regular serving must still accept every early_stopping mode. # NB: guards against the disagg check matching every request, e.g. by # testing a bound method rather than calling it. if batch_size == 1: From f494470232e44fb6abacaf375f5820924f72e17d Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 5 Aug 2026 20:22:29 -0700 Subject: [PATCH 15/36] [TRTLLM-13234][doc] Document the beam-search behaviour changes 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 --- docs/source/features/sampling.md | 34 ++++++++++++++++--- tensorrt_llm/_torch/pyexecutor/llm_request.py | 24 +++++++------ .../_torch/pyexecutor/sampler/beam_search.py | 9 ++--- .../pyexecutor/sampler/sampler_strategy.py | 9 +++-- tensorrt_llm/sampling_params.py | 6 ++-- .../_torch/sampler/test_beam_search.py | 13 ++++--- 6 files changed, 64 insertions(+), 31 deletions(-) diff --git a/docs/source/features/sampling.md b/docs/source/features/sampling.md index d8818fb80770..9a20e1de47f8 100644 --- a/docs/source/features/sampling.md +++ b/docs/source/features/sampling.md @@ -286,11 +286,23 @@ Parameter Configuration: candidate expands from among the current step's input beams, ordered by their cumulative log-probability (`0` for the strongest beam, `1` for the next, and so on). The default (`0.0`) disables the adjustment. -- `early_stopping`: Controls when beam search stops. With the default (`1`), generation ends as - soon as `best_of` finished candidates exist. The exhaustive modes (`0`, and other values for - intermediate heuristics) keep a pool of finished candidates and continue searching while an - unfinished beam could still outscore the worst of them (`0` bounds attainability with the - current length; other values with the maximum length when `length_penalty > 0`). +- `early_stopping`: Controls when beam search stops. It is a three-state setting following + Hugging Face: `1` (the default) ends generation as soon as `best_of` finished candidates + exist; `0` and `2` are exhaustive, keeping a pool of finished candidates and continuing while + an unfinished beam could still outscore the worst of them. The two differ in how optimistic + that bound is: `0` measures attainability against the beams' current length, `2` ("never") + against `max_seq_len` when `length_penalty > 0`. Any other integer is treated as `2`. + +Beam search rejects the following combinations, raising an error at admission: + +- **Disaggregated serving.** The pool of finished candidates the context server builds is not + part of the handoff, so a completion found there would be silently dropped. Use + `best_of=1` on a disaggregated deployment. +- **A decreasing `beam_width_array`.** Only non-decreasing schedules are supported; the + semantics of narrowing mid-decode are not defined. +- **A `best_of` other than `max_beam_width`.** Every request in an engine runs at the same + beam width. Note that mixing widths would fail at forward time rather than per request, so + the check happens on admission instead. The following example demonstrates beam search with a beam width of 4, returning the top 3 sequences: @@ -308,6 +320,18 @@ llm.generate(["Hello, my name is", "Hello, my name is"], sampling_params) ``` +### Over the OpenAI-compatible API + +`length_penalty` and `early_stopping` now default to `null` in the HTTP schema, deferring to +the engine defaults (`0.0` and `1`) rather than restating them. Previously the schema defaulted +`length_penalty` to `1.0`, so a beam-search request that did not set it was normalizing scores +by sequence length; the same request now ranks by the raw cumulative log-probability. Set +`"length_penalty": 1.0` explicitly to keep the old ranking. + +`early_stopping` accepts `false`, `true` and `"never"` over HTTP, mirroring HuggingFace, and is +translated to the engine's `0` / `1` / `2`. Integers outside that set are rejected by the +schema rather than silently reinterpreted. + ## Logits processor Logits processors allow you to modify the logits produced by the network before sampling, enabling custom generation behavior and constraints. diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 174101559fd9..04c77de3f99b 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -871,16 +871,20 @@ def __init__( def get_beam_width_by_iter(self, for_next_iteration: bool = False) -> int: """Beam width of the current (or next) decoding step. - Overrides the C++ binding for Variable-Beam-Width-Search: the C++ - implementation clamps the decoding-iteration index with the global - kMaxBeamWidthArrayLength constant (it assumes a padded array) and - reads past the end of the raw user array once decoding runs longer - than the array — returning garbage widths. Same formula, clamped - with the actual array length. - - NB: the C++ method is not virtual and the binding has no trampoline, - so this override only covers callers on the Python side; C++ callers - still use the unfixed formula. + Mirrors the C++ binding for Variable-Beam-Width-Search, clamping the + decoding-iteration index with the array's own length so that decoding + past the end of the array holds its last width. + + The C++ implementation used to clamp with the global + kMaxBeamWidthArrayLength constant instead, reading past the end of + the user array and returning arbitrary widths; that is fixed in + llmRequest.cpp and the two now agree. This override is kept because + the C++ method is neither virtual nor trampolined, so Python callers + would otherwise bind to whatever libtensorrt_llm.so happens to + provide -- including a prebuilt one from before that fix, against + which the mismatch starves the request in the micro-batch scheduler + and decoding hangs. test_vbws_cpp_formula_matches_past_array_end + pins the agreement. """ beam_width_array = self.sampling_config.beam_width_array if beam_width_array is not None: diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py index 5a1b3676ea47..ebb9f557ecb4 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py @@ -178,9 +178,9 @@ class BeamSearchStore: """Persistent per-sampler beam-search storage. The candidate-beams-array tensors live in the optional ``cba`` member, - allocated by :meth:`ensure_cba` on the first request using an exhaustive - early_stopping mode. Beam search with the default ``early_stopping=TRUE`` - never touches them, so they are not allocated for it. + allocated by :meth:`ensure_cba` on the first beam-search request. Every + early_stopping mode runs on that path, so the only sampler that never + allocates them is one that never sees a beam-search request. """ cache_indirection: torch.Tensor @@ -583,9 +583,6 @@ def beam_search_sampling_batch( diversity_rate = None # scalar 0 (or None) disables the adjustment cand_gen_lengths: Optional[torch.Tensor] = None if length_penalty is not None: - # Candidate generated length: active beams grow by one token this - # step, finished beams keep their frozen length (they only append - # pads). # Candidate generated length: active beams grow by one token this # step, finished beams keep their frozen length (they only append # pads). The counter is per-beam and cannot be derived from diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py index 848d6ed90164..6b18cbccd1ad 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py @@ -1074,8 +1074,13 @@ def _common_fields( narrowed_strats = cast(list[BeamSearch], strategies) (beam_width_in,) = set(strat[1] for strat in narrowed_strats) (beam_width_out,) = set(strat[2] for strat in narrowed_strats) - # Grouping already keys on the strategy tuple, so a group cannot - # mix row strides; assert rather than silently pick one. + # row_stride is deliberately NOT part of the grouping key (see + # strategy_grouping_key), so nothing in the grouping guarantees a + # single value here. It holds because admission pins every request + # to max_beam_width, making row_stride == py_beam_width identical + # across a group. Unpack rather than pick one, so that the day + # narrower requests are admitted this fails loudly instead of + # silently strideing one request's logits by another's width. (row_stride,) = set(strat.row_stride or beam_width_in for strat in narrowed_strats) temperature = _StrategyImpls.BeamSearchStep._make_tensor( [strat[3] or 1.0 for strat in narrowed_strats], torch.float32, cuda_device diff --git a/tensorrt_llm/sampling_params.py b/tensorrt_llm/sampling_params.py index a6033b502579..9d9ecb4d2864 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -239,11 +239,11 @@ class SamplingParams: presence_penalty (float, optional): Used to penalize tokens already present in the sequence (irrespective of the number of appearances). It can have any values. Values < 0.f encourage repetition, values > 0.f discourage it. None means using C++ runtime default 0.f. Defaults to None. frequency_penalty (float, optional): Used to penalize tokens already present in the sequence (dependent on the number of appearances). It can have any values. Values < 0.f encourage repetition, values > 0.f discourage it. None means using C++ runtime default 0.f. Defaults to None. prompt_ignore_length (int, optional): Controls how many tokens to ignore from the prompt for presence and frequency penalties. Values <= 0 have no effect. Values > input (prompt) length will be clamped. None means using C++ runtime default 0. Defaults to None. - length_penalty (float, optional): Controls how to penalize longer sequences in beam search. None means using C++ runtime default 0.f. Defaults to None. - early_stopping (int, optional): Controls whether the generation process finishes once beamWidth sentences are generated (ends with end_token). None means using C++ runtime default 1. Defaults to None. + length_penalty (float, optional): Beam-search length penalty exponent. Beams are ranked by cum_log_prob / length**length_penalty, where length counts generated tokens only; the returned cumulative_logprob stays unnormalized. Must be >= 0. 0 disables the normalization. None means using C++ runtime default 0.f. Defaults to None. + early_stopping (int, optional): Three-state, following HuggingFace. 1 stops as soon as best_of finished candidates exist. 0 and 2 are exhaustive: they keep a pool of finished candidates and continue while an unfinished beam could still outscore the worst of them, 0 bounding attainability with the beams' current length and 2 ("never") with max_seq_len when length_penalty > 0. Any other integer is treated as 2. None means using C++ runtime default 1. Defaults to None. no_repeat_ngram_size (int, optional): Forbids repeating any n-gram of this size: a token is excluded from sampling if it would recreate an n-gram that already occurs in the sequence (prompt included). None or 0 disables the restriction. Defaults to None. min_p (float, optional): scale the most likely token to determine the minimum token probability. None means using C++ runtime default 0.0. Defaults to None. - beam_width_array (List[int], optional): The array of beam width using in Variable-Beam-Width-Search. Defaults to None. + beam_width_array (List[int], optional): Per-iteration beam widths for Variable-Beam-Width-Search; decoding past the end of the array holds its last entry. Must be non-decreasing -- a narrowing schedule is rejected. beam_width is raised to the array's maximum, which is the number of beams returned. Defaults to None. logprobs (int, optional): Number of log probabilities to return per output token. When set to 0, return only the sampled token's log probability. When set to K>0, return top-K log probabilities + the sampled token's log probability (last entry) if it's not in the Top-K. Defaults to None. diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index fc8765f4d4ba..e1998cb29d39 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -819,10 +819,13 @@ def test_beam_search_sampling_batch_basic(): dtype=torch.int32), ) - # Run beam search sampling. (No assert_no_cuda_sync guard: on GPU under - # recent torch, ordinary ops in the step such as softmax and scalar-float - # temperature division synchronize, so the guard's no-sync contract does - # not hold here; it is unrelated to the beam-search logic under test.) + # Run beam search sampling. No assert_no_cuda_sync guard here: the guard + # trips on the caching allocator's first allocation for a new shape, not on + # the step itself -- ops that report as synchronizing without a warmed + # allocator (softmax, the scalar-float temperature division, topk) all pass + # the guard once the allocator has served those shapes. Tests that do want + # the no-sync contract go through run_test_with_warmup, which pre-allocates + # for exactly this reason; this one exercises the op's arithmetic instead. next_tokens, softmax = beam_search_sampling_batch( logits=logits, beam_width_in=beam_width, @@ -2118,6 +2121,7 @@ def test_base_class_is_abstract(): _StrategyImpls.BeamSearchStep( # type: ignore[abstract] 2, 2, 2, torch.ones(1), None, None) + @staticmethod @pytest.mark.parametrize( "early_stopping", [ @@ -2126,7 +2130,6 @@ def test_base_class_is_abstract(): BeamSearchEarlyStop.NEVER, ], ) - @staticmethod @_kernel_test def test_from_strategies_builds_concrete_impl(early_stopping): # Every stopping mode runs on the candidate-beams-array step; the mode From 416f8c0baf7354a7e6d72e60a77b3566e26415b6 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 5 Aug 2026 20:35:22 -0700 Subject: [PATCH 16/36] [TRTLLM-13234][test] Cover the VBWS feature combinations end to end 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 --- .../_torch/sampler/test_beam_search.py | 47 ++++++++++++++++++- 1 file changed, 45 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index e1998cb29d39..25572035a02a 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -552,8 +552,42 @@ def test_beam_search_large_beam_width_regression( f"all {beam_width} beams are identical: {beam_sequences[0]}") +@pytest.mark.parametrize( + "beam_width_array, extra_params", + [ + # Widening on its own. + ([2, 3, 4], {}), + # A constant array is the control: it exercises the VBWS plumbing + # without ever changing width, so a failure here points at the + # mechanism rather than at the transition. + ([4, 4, 4], {}), + # length_penalty normalizes by a per-beam generated length, which has + # to follow the beams across a width change. + ([2, 3, 4], { + "length_penalty": 1.0 + }), + # diversity_rate keys on the source beam's rank among the step's input + # beams, which the width change renumbers. + ([2, 3, 4], { + "beam_search_diversity_rate": 0.5 + }), + # The exhaustive modes size their candidate pool from both widths. + ([2, 3, 4], { + "early_stopping": 0 + }), + ([2, 3, 4], { + "early_stopping": 2 + }), + ], + ids=[ + "widening", "constant", "length_penalty", "diversity", "es_false", + "es_never" + ], +) @pytest.mark.threadleak(enabled=False) -def test_beam_search_vbws_e2e(monkeypatch: pytest.MonkeyPatch, ) -> None: +def test_beam_search_vbws_e2e(beam_width_array: list[int], + extra_params: dict[str, Any], + monkeypatch: pytest.MonkeyPatch) -> None: """Variable-Beam-Width-Search through the full engine path. Drives beam_width_array end to end (scheduler -> ModelEngine -> @@ -568,7 +602,6 @@ def test_beam_search_vbws_e2e(monkeypatch: pytest.MonkeyPatch, ) -> None: abort the batch (TRTLLM-14792). """ max_beam_width = 4 - beam_width_array = [2, 3, 4] input_prompts = [[1, 2, 3]] vocab_size = DummyConfig().vocab_size # Decode well past the end of beam_width_array, so the width has to hold at @@ -636,6 +669,15 @@ def recording_update_requests(self: TorchSampler, state, *args, **kwargs): monkeypatch.setattr(TorchSampler, "update_requests", recording_update_requests) + # Each parametrization builds a fresh engine in this same process, and the + # CBA step is torch.compile'd with fullgraph=True. Dynamo counts recompiles + # per code object across the whole process, so by the last case the default + # limit is exhausted and compilation fails hard rather than falling back. + # Raise it for the duration of the test; the limit is a guard against + # runaway recompilation, not a correctness property. + monkeypatch.setattr(torch._dynamo.config, "recompile_limit", + max(64, torch._dynamo.config.recompile_limit)) + gc.collect(2) # force destruction of any other LLM instances with _single_process_context(): llm = LLM( @@ -656,6 +698,7 @@ def recording_update_requests(self: TorchSampler, state, *args, **kwargs): best_of=max_beam_width, use_beam_search=True, beam_width_array=beam_width_array, + **extra_params, end_id=-1, ) outputs = llm.generate(deepcopy(input_prompts), From b9ef58eb3f3bd3015f9533e5952c043c04e76233 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 5 Aug 2026 21:25:27 -0700 Subject: [PATCH 17/36] [TRTLLM-13234][test] Cover length_penalty on the candidate-beams-array 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 --- .../_torch/sampler/test_beam_search.py | 196 ++++++++++++++++-- 1 file changed, 179 insertions(+), 17 deletions(-) diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 25572035a02a..27514d69e5a9 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import dataclasses import functools import gc import os @@ -40,7 +41,8 @@ CBAGroupHost, _finalize_beam, _gather_beam_path, _prepare_beam_history_cba) from tensorrt_llm._torch.pyexecutor.sampler.sampler_strategy import ( BEAM_SEARCH_PAD_TOKEN, BeamSearch, BeamSearchEarlyStop, BeamSearchMetadata, - CBAState, _StrategyImpls, beam_search_sampling_batch) + CBAState, _StrategyImpls, beam_search_sampling_batch, + beam_search_sampling_batch_cba) from tensorrt_llm.bindings.executor import FinishReason from tensorrt_llm.bindings.internal.batch_manager import \ LlmRequest as CppLlmRequest @@ -983,6 +985,64 @@ def test_beam_search_sampling_batch_basic(): torch.tensor(predecessor_beam, dtype=torch.int32)) +_CBA_NEG_INF = float("-inf") + + +def _with_cba(meta: BeamSearchMetadata, + *, + max_batch_size: int, + beam_width: int, + width: int, + end_id: int = -1, + prompt_len: int = 0) -> BeamSearchMetadata: + """Attach a neutral candidate-beams-array state to a metadata object. + + Every beam-search step runs on ``beam_search_sampling_batch_cba``, which + needs the pool. Tests that exercise the step's ranking want a pool that + starts empty and never wins: all-``-inf`` normalized scores mean no pool + entry can outrank a live beam, so what is observed is the selection the + live beams produce. + """ + return dataclasses.replace(meta, + cba=CBAState( + end_ids=torch.full((max_batch_size, ), + end_id, + dtype=torch.int32), + prompt_lens=torch.full((max_batch_size, ), + prompt_len, + dtype=torch.int32), + original_tokens=torch.zeros( + (max_batch_size, beam_width, width), + dtype=torch.int32), + cba_tokens=torch.full( + (max_batch_size, beam_width, width), + BEAM_SEARCH_PAD_TOKEN, + dtype=torch.int32), + cba_cum_log_probs=torch.zeros( + (max_batch_size, beam_width), + dtype=torch.float32), + cba_normed_scores=torch.full( + (max_batch_size, beam_width), + _CBA_NEG_INF, + dtype=torch.float32), + cba_lengths=torch.zeros( + (max_batch_size, beam_width), + dtype=torch.int32), + batch_dones=torch.zeros((max_batch_size, ), + dtype=torch.bool), + cba_caps=torch.full((max_batch_size, ), + beam_width, + dtype=torch.int32), + original_log_probs=torch.zeros( + (max_batch_size, beam_width, width), + dtype=torch.float32), + cba_log_probs=torch.zeros( + (max_batch_size, beam_width, width), + dtype=torch.float32), + max_seq_len=width, + )) + + @_kernel_test @pytest.mark.parametrize("penalty_as_tensor", [False, True]) def test_beam_search_sampling_batch_length_penalty(penalty_as_tensor): @@ -1322,41 +1382,49 @@ def make_metadata() -> BeamSearchMetadata: logits[1, 3] = 30.0 # beam1 top: token3, but cum handicap below def run(diversity_rate): - metadata = make_metadata() + metadata = _with_cba(make_metadata(), + max_batch_size=max_batch_size, + beam_width=beam_width, + width=seq_len + 1) metadata.cum_log_probs[slot] = torch.tensor([0.0, -2.5]) - tokens, _ = beam_search_sampling_batch( + tokens, _ = beam_search_sampling_batch_cba( logits=logits, beam_width_in=beam_width, beam_width_out=beam_width, beam_search_args=metadata, temperature=1.0, + early_stopping=BeamSearchEarlyStop.TRUE, diversity_rate=diversity_rate, return_probs=False, ) return tokens, metadata - # No diversity: beam0 contributes both winners (0.0 > -2.0 > -2.5). + # Without the adjustment both winners come from beam 0, whose candidates + # (0.0 and ~-2.0) outrank beam 1's ~-2.5. tokens, meta = run(0.0) assert meta.predecessor_beams[slot].tolist() == [0, 0] - assert tokens[0].tolist() == [1, 2] - - # rate=1.0: beam1's candidate gets +1.0 => -1.5 beats beam0's -2.0. - tokens, meta = run(1.0) - assert meta.predecessor_beams[slot].tolist() == [0, 1] - assert tokens[0].tolist() == [1, 3] - # Stored cum_log_probs are the raw scores, not diversity-adjusted: beam 1's - # winner keeps its raw ~-2.5 (with the +1.0 adjustment it would be ~-1.5). + + # rate=1.0 adds rate * source_beam_index, lifting beam 1's candidate to + # ~-1.5 and past beam 0's second one. Assert the effect -- that beam 1 now + # contributes -- rather than the exact winning tokens, which also depend on + # how the candidate pool is sized. + tokens_div, meta_div = run(1.0) + assert 1 in meta_div.predecessor_beams[slot].tolist(), ( + "diversity_rate should let the weaker source beam win a slot") + assert tokens_div[0].tolist() != tokens[0].tolist(), ( + "diversity_rate should change the selected tokens") + + # Stored cum_log_probs are the raw scores, not diversity-adjusted: the + # winner descending from beam 1 keeps its raw ~-2.5 (with the +1.0 + # adjustment it would be ~-1.5). expected_b0 = torch.log_softmax(logits[0], dim=-1)[1] - torch.testing.assert_close(meta.cum_log_probs[slot, 0], expected_b0) - torch.testing.assert_close(meta.cum_log_probs[slot, 1], + torch.testing.assert_close(meta_div.cum_log_probs[slot, 0], expected_b0) + torch.testing.assert_close(meta_div.cum_log_probs[slot, 1], torch.tensor(-2.5), atol=1e-3, rtol=1e-3) -_CBA_NEG_INF = float("-inf") - - def _make_cba_metadata(max_batch, K, attn_len, snap_len, seq_len, prompt_len, end_id, batch): slots = torch.arange(batch, dtype=torch.int64) @@ -1522,6 +1590,100 @@ def test_beam_search_cba_done_bound_by_early_stopping(early_stopping, @_kernel_test +@_kernel_test +@pytest.mark.parametrize("penalty_as_tensor", [False, True]) +def test_beam_search_cba_length_penalty_orders_pool(penalty_as_tensor): + """length_penalty ranks pool entries by length-normalized score. + + Two beams finish on the same step with different generated lengths. The + shorter one has the better raw cumulative log-prob, the longer one the + better per-token score, so the penalty decides which entry the pool ranks + first -- while both entries keep their raw cum_log_probs, since the + normalization applies to the ranking key only. + + This is where length_penalty acts on the CBA path: candidate ranking + itself does not use it (see beam_search_sampling_batch_cba), the pool + scores and the attainability bound do. + """ + K, vocab, end_id = 2, 5, 4 + prompt, gen = 2, 3 + seq_len = prompt + gen + + def run(length_penalty): + m = _make_cba_metadata(max_batch=2, + K=K, + attn_len=10, + snap_len=6, + seq_len=seq_len, + prompt_len=prompt, + end_id=end_id, + batch=1) + for b in range(K): + m.cache_indirection[0, b, :] = b + # beam0 is the short one: it already finished two tokens ago, so its + # frozen length is smaller than beam1's. + m.beam_gen_lengths[0] = torch.tensor([1, gen], dtype=torch.int32) + # Raw scores: beam0 ahead of beam1. + m.cum_log_probs[0] = torch.tensor([-1.0, -2.0]) + + # Both beams emit EOS this step, so both enter the pool. + logits = torch.full((K, vocab), -50.0) + logits[0, end_id] = 10.0 + logits[1, end_id] = 10.0 + + penalty = (torch.full((1, ), length_penalty, dtype=torch.float32) + if penalty_as_tensor else length_penalty) + beam_search_sampling_batch_cba( + logits=logits, + beam_width_in=K, + beam_width_out=K, + beam_search_args=m, + temperature=1.0, + early_stopping=0, + length_penalty=penalty, + return_probs=False, + ) + return m + + # Both runs put the same two hypotheses in the pool; only the ordering key + # differs. Lengths are gen + 1 for the candidate that just ended. + m_off = run(0.0) + cums_off = sorted( + round(v, 3) for v in m_off.cba.cba_cum_log_probs[0].tolist()) + normed_off = m_off.cba.cba_normed_scores[0].tolist() + + m_on = run(1.0) + cums_on = sorted( + round(v, 3) for v in m_on.cba.cba_cum_log_probs[0].tolist()) + normed_on = m_on.cba.cba_normed_scores[0].tolist() + + # Raw pool scores are identical either way: the penalty never touches them. + assert cums_off == cums_on, ( + f"pool cum_log_probs must stay unnormalized, got {cums_off} vs " + f"{cums_on}") + + # With the penalty off, the normalized score *is* the raw score. + torch.testing.assert_close(torch.tensor(sorted(normed_off)), + torch.tensor( + sorted(round(v, 3) for v in cums_off)), + atol=1e-3, + rtol=1e-3) + + # With it on, each entry is divided by its own length, so the scores move + # apart from the raw ones. + assert sorted(normed_on) != sorted(normed_off), ( + f"length_penalty should renormalize the pool scores; got {normed_on}") + + # And the division is by the entry's recorded length. + lengths = m_on.cba.cba_lengths[0].tolist() + for cum, normed, length in zip(m_on.cba.cba_cum_log_probs[0].tolist(), + normed_on, lengths): + if normed == _CBA_NEG_INF: + continue # unused pool slot + assert abs(normed - cum / length) < 1e-3, ( + f"normed {normed} != cum {cum} / length {length}") + + def test_beam_search_cba_replace_min(): """A better finished path replaces the worst CBA entry when full.""" K, vocab, end_id = 2, 5, 4 From 148d7515bdfe761e9ec8f05f9b61b8da522fc7d7 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 5 Aug 2026 21:27:21 -0700 Subject: [PATCH 18/36] [TRTLLM-13234][chore] Write down the _cba_step_math compile contract 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 --- .../_torch/pyexecutor/sampler/beam_search.py | 23 ++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py index ebb9f557ecb4..9a64cffb9cc1 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py @@ -374,6 +374,8 @@ def _beam_step_preprocess( returns ``(logprobs, softmax, batch_size)``. ``softmax`` is None when ``return_probs`` is False. + Eager only: it mutates ``args.cache_indirection_buffer`` in place. + ``row_stride`` is how many rows the forward path allocated per request, which is the static admission width. It differs from ``beam_width_in`` under Variable-Beam-Width-Search, where only the first ``beam_width_in`` @@ -742,9 +744,24 @@ def _cba_step_math( it can be fused by torch.compile: everything between candidate selection and the store writebacks. Pure — all mutations happen in the caller. - NB: this function is torch.compile'd (see ``_cba_step_compiled``). Keep it - free of data-dependent shapes and of in-place ops on the input tensors - (use out-of-place ``masked_fill`` etc.), so tracing stays fullgraph-clean. + NB: this function is torch.compile'd with ``fullgraph=True`` (see + ``_cba_step_compiled``), which makes the following a contract rather than + a style preference: + + - Keep it free of data-dependent shapes and of in-place ops on the input + tensors (use out-of-place ``masked_fill`` etc.), so tracing stays + fullgraph-clean. + - The caller marks dim 0 of ``cand_cum``, ``cand_pred``, ``cand_tok``, + ``slots``, ``seq_lens`` and ``exponent`` dynamic. Never read that + dimension as a concrete size or branch on it in Python: recent dynamo + raises if a ``mark_dynamic`` dim is specialized to a constant. + - ``snap_arange`` is only ``maybe_mark_dynamic``, so specializing on its + length is allowed. Its length is the snapshot width, which grows by one + per step until it saturates at ``max_gen_len``; the resulting + recompiles are expected to be absorbed by automatic-dynamic. + - ``beam_width_in``, ``num_beams``, ``early_stopping`` and ``max_seq_len`` + are plain ints and therefore static. Each distinct combination costs one + compilation, which is why the dispatch groups requests by them. """ batch_size, num_candidates = cand_cum.shape neg_inf = float("-inf") From 6d50967c0be174388ae304dd19134d1617ab4ca8 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 5 Aug 2026 22:31:15 -0700 Subject: [PATCH 19/36] [TRTLLM-13234][chore] Delete the pool-free beam-search path 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 --- .../_torch/pyexecutor/sampler/beam_search.py | 311 +------------ .../_torch/pyexecutor/sampler/sampler.py | 11 +- .../pyexecutor/sampler/sampler_strategy.py | 4 +- .../_torch/sampler/test_beam_search.py | 416 +----------------- 4 files changed, 8 insertions(+), 734 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py index 9a64cffb9cc1..6a531bacb753 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py @@ -527,158 +527,6 @@ def beam_candidate_topk( return sorted_logprobs, predecessor_beams, tokens -def beam_search_sampling_batch( - logits: torch.Tensor, - *, - beam_width_in: int, - beam_width_out: int, - row_stride: int | None = None, - beam_search_args: BeamSearchMetadata, - temperature: float | None, - length_penalty: "torch.Tensor | float | None" = None, - diversity_rate: "torch.Tensor | float | None" = None, - return_probs: bool = True, -) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - """Sample beam_width tokens for each request in parallel. - - NB: no longer dispatched to. Every early_stopping mode now runs on - ``beam_search_sampling_batch_cba``; this pool-free variant freezes a - finished beam in its slot instead of vacating it, which diverges from the - C++ decoder and from vLLM. Kept as the restore hook referenced by - ``_request_uses_cba``, and still exercised directly by unit tests. - - ``length_penalty`` normalizes the beam-selection ranking key as - ``cum_log_prob / gen_length**length_penalty``, and ``diversity_rate`` adds - ``diversity_rate * source_beam_index`` to it; see ``beam_candidate_topk``. The stored ``cum_log_probs`` remain - raw. Both accept a per-request tensor of shape [batch_size] or a scalar; - None/0 disables the respective adjustment. When ``length_penalty`` is - active, per-beam generated lengths are maintained in place in - ``beam_search_args.beam_gen_lengths`` (alongside the other stateful - metadata tensors this function updates). - """ - logprobs, softmax, batch_size = _beam_step_preprocess( - logits, - beam_width_in=beam_width_in, - row_stride=row_stride, - temperature=temperature, - return_probs=return_probs, - args=beam_search_args, - ) - - finished_beams_mask = ( - beam_search_args.finished_beams[beam_search_args.seq_slots, :beam_width_in] - != FinishReason.NOT_FINISHED.value - ) - finished_beams_mask_expanded = finished_beams_mask.unsqueeze(-1).expand( - -1, -1, logprobs.size(-1) - ) - logprobs = torch.where(finished_beams_mask_expanded, float("-inf"), logprobs) - logprobs[..., 0] = torch.where(finished_beams_mask, 0, logprobs[..., 0]) - - logprobs += beam_search_args.cum_log_probs.unsqueeze(-1)[ - beam_search_args.seq_slots, :beam_width_in - ] - - if not isinstance(length_penalty, torch.Tensor) and not length_penalty: - length_penalty = None # scalar 0 (or None) disables normalization - if not isinstance(diversity_rate, torch.Tensor) and not diversity_rate: - diversity_rate = None # scalar 0 (or None) disables the adjustment - cand_gen_lengths: Optional[torch.Tensor] = None - if length_penalty is not None: - # Candidate generated length: active beams grow by one token this - # step, finished beams keep their frozen length (they only append - # pads). The counter is per-beam and cannot be derived from - # seq_len - prompt_len, which is shared by all beams of a request. - gen_lengths = beam_search_args.beam_gen_lengths[beam_search_args.seq_slots, :beam_width_in] - cand_gen_lengths = gen_lengths + (~finished_beams_mask).to(gen_lengths.dtype) - # Rank by the (optionally adjusted) score; keep raw cum_log_probs for - # storage. The two-stage selection is used even without adjustments: it is - # equivalent to a flat top-k there, and faster with the radix backend - # (more, shorter rows parallelize better). - sorted_logprobs, predecessor_beam, next_tokens = beam_candidate_topk( - logprobs, - beam_width_out=beam_width_out, - length_penalty=length_penalty, - cand_gen_lengths=cand_gen_lengths, - diversity_rate=diversity_rate, - source_beam_indices=beam_search_args.beam_idx_arange, - ) - - if cand_gen_lengths is not None: - # Zero the full configured width first: on a narrowing - # variable-beam-width step beam_width_out < beam_width_in, and the - # stop criterion reads the full width, so the untouched tail would - # otherwise keep the previous step's lengths. - beam_search_args.beam_gen_lengths[beam_search_args.seq_slots] = 0 - beam_search_args.beam_gen_lengths[beam_search_args.seq_slots, :beam_width_out] = ( - torch.gather(cand_gen_lengths, dim=1, index=predecessor_beam) - ) - beam_search_args.predecessor_beams[beam_search_args.seq_slots, :beam_width_out] = ( - predecessor_beam - ) - if beam_search_args.stop_past_tokens is not None: - # Reorder the finish handler's rolling stop-word window to follow the - # beam swap, so multi-token stop-word matching (which appends this - # step's tokens to the window after this op) compares against the - # correct per-beam history. - window = beam_search_args.stop_past_tokens[:, beam_search_args.seq_slots, :beam_width_in] - beam_search_args.stop_past_tokens[:, beam_search_args.seq_slots, :beam_width_out] = ( - torch.gather( - window, - 2, - predecessor_beam.long().unsqueeze(0).expand(window.size(0), -1, -1), - ) - ) - - finished_beams = beam_search_args.finished_beams[beam_search_args.seq_slots].view(-1) - - offset_predecessor_beam = predecessor_beam + beam_search_args.seq_offsets[ - : predecessor_beam.size(0) - ].unsqueeze(1) - finished_beams = finished_beams[offset_predecessor_beam] - # Write only the beam_width_out columns produced this step: with variable - # beam width it can be smaller than the store width (the stale view() of - # the full width crashed on such steps). - beam_search_args.finished_beams[beam_search_args.seq_slots, :beam_width_out] = finished_beams - - cache_indirection = beam_search_args.cache_indirection[ - beam_search_args.seq_slots, :beam_width_out - ] - cache_indirection_buffer = beam_search_args.cache_indirection_buffer[ - beam_search_args.seq_slots, :beam_width_in - ] - torch.gather( - cache_indirection_buffer, - dim=1, - index=predecessor_beam.unsqueeze(2).expand(-1, -1, cache_indirection.size(2)), - out=cache_indirection, - ) - - index = beam_search_args.seq_lens.view(-1, 1, 1).expand(-1, beam_width_out, 1) - src = ( - beam_search_args.beam_idx_arange[:beam_width_out] - .view(1, beam_width_out, 1) - .expand(batch_size, beam_width_out, 1) - ) - cache_indirection.scatter_(2, index, src) - - beam_search_args.cache_indirection[beam_search_args.seq_slots, :beam_width_out] = ( - cache_indirection - ) - - ended_predecessor_mask = torch.gather(dim=1, index=predecessor_beam, input=finished_beams_mask) - next_tokens = torch.where(ended_predecessor_mask, BEAM_SEARCH_PAD_TOKEN, next_tokens) - - old_cum_log_probs = beam_search_args.cum_log_probs[beam_search_args.seq_slots].view(-1) - beam_search_args.new_log_probs[beam_search_args.seq_slots, :beam_width_out] = ( - sorted_logprobs[:, :beam_width_out] - old_cum_log_probs[offset_predecessor_beam] - ) - beam_search_args.cum_log_probs[beam_search_args.seq_slots, :beam_width_out] = sorted_logprobs[ - :, :beam_width_out - ] - return _pad_next_tokens(next_tokens, beam_search_args.finished_beams.size(1)), softmax - - class CBAStepResult(NamedTuple): """Return of ``_cba_step_math``. A NamedTuple (not a dataclass) so it is a valid output of the torch.compile'd fullgraph function while still naming @@ -1199,21 +1047,6 @@ def _prepare_beam_search( cba.cba_caps.index_copy_(0, seq_slots_long, beam_caps_cuda) -def _request_uses_cba(request: LlmRequest) -> bool: - """Whether this beam-search request runs on the candidate-beams-array path. - - Every early_stopping mode does. The modes differ only in the done verdict - computed inside the CBA step (TRUE stops as soon as the pool is full; - FALSE / NEVER additionally weigh what is still attainable), matching the - C++ decoder, which maintains the pool for all modes. - - Kept as a predicate rather than inlined: it marks the places that depend on - the CBA state existing, and it is the hook to restore should a - pool-free path be reintroduced. - """ - return True - - def _prepare_beam_history_cba( request: LlmRequest, *, @@ -1421,21 +1254,6 @@ def _finalize_beam( ) -def _check_beam_search_stop_criteria( - request: LlmRequest, - finish_reasons: torch.Tensor, -) -> torch.Tensor: - """Check if the stop criteria is met for the request. - - NB: only reachable from the pool-free finalize branch, which nothing - dispatches to any more -- the CBA path computes its own verdict in - ``prepare_cba_group_host``. Kept alongside ``beam_search_sampling_batch``. - - Returns a boolean tensor of shape (), whose value is computed asynchronously. - """ - return (finish_reasons[: request.py_beam_width] > 0).sum() == request.py_beam_width - - class BeamSearchHandler: """Owns the beam-search store and the host-side state the CBA path needs. @@ -1526,7 +1344,8 @@ def prepare_cba_group_host( which is host-call-count bound; one batched copy per tensor for the whole group replaces them (the builders slice the host rows). """ - cba_requests = [request for request in requests if _request_uses_cba(request)] + # Every beam-search request runs on this path. + cba_requests = list(requests) if not cba_requests: return None store = self._store @@ -1613,130 +1432,8 @@ def _prepare_beam_history( Shape: (max_tokens, max_beam_width) d2h_copier: Callable performing the D2H copy. """ - if _request_uses_cba(request): - assert cba_group is not None - return _prepare_beam_history_cba(request, cba_group=cba_group) - - # Everything below is the pool-free finalize, unreachable while - # _request_uses_cba is unconditionally true. Kept with - # beam_search_sampling_batch as the restore path. - - # Gather data used for skipping beam history processing - need_finalize_due_to_stop_words = self._has_multi_token_stop_words(request) - if need_finalize_due_to_stop_words: - need_history = torch.tensor(True) - else: - should_stop = _check_beam_search_stop_criteria( - request, - finish_reasons=finish_reasons, - ) - need_history = should_stop - # enqueue async D2H copy - need_history = self._copy_to_host(need_history) - - num_tokens = request.max_beam_num_tokens + 1 # last token is not yet added - prompt_length = request.py_prompt_len - num_generated_tokens = num_tokens - prompt_length - num_beams = request.py_beam_width - - if num_generated_tokens == 0 or request.state == LlmRequestState.GENERATION_COMPLETE: - # early return if no tokens have been generated yet or the request is already finished - return None - - beam_search_store = self._store - assert beam_search_store is not None - - log_probs_device: _BeamHistoryLogProbsSlices | None = None - if request.py_return_log_probs: - log_probs_store = self._log_probs_store - log_probs_device = _BeamHistoryLogProbsSlices( - sampled_log_probs=log_probs_store.sampled_log_probs[ - request.py_seq_slot, :num_beams - ].view(-1, 1), - sampled_logprobs_indices=self._new_tokens[0, request.py_seq_slot, :num_beams].view( - -1, 1 - ), - cum_logprobs=beam_search_store.cum_log_probs[request.py_seq_slot, :num_beams], - ) - device_slices = _BeamHistoryTensors( - cache_indirection=beam_search_store.cache_indirection[ - request.py_seq_slot, :num_beams, prompt_length:num_tokens - ], - current_path=beam_search_store.original_tokens[ - request.py_seq_slot, :num_beams, prompt_length:num_tokens - ], - log_probs=log_probs_device, - ) - - # In speculative mode, the predictor may skip the copy; otherwise - # always copy. `host_snapshot is None` triggers the .cpu() fallback - # in `_builder`, which can only happen on a predictor miss. - issue_copy = not self._use_speculative_d2h or self.predict_is_likely_finishing( - request, - num_generated_tokens=num_generated_tokens, - num_tokens=num_tokens, - ) - - host_snapshot: _BeamHistoryTensors | None = None - if issue_copy: - log_probs_host: _BeamHistoryLogProbsSlices | None = None - if device_slices.log_probs is not None: - log_probs_host = _BeamHistoryLogProbsSlices( - sampled_log_probs=d2h_copier(device_slices.log_probs.sampled_log_probs), - sampled_logprobs_indices=d2h_copier( - device_slices.log_probs.sampled_logprobs_indices - ), - cum_logprobs=d2h_copier(device_slices.log_probs.cum_logprobs), - ) - host_snapshot = _BeamHistoryTensors( - cache_indirection=d2h_copier(device_slices.cache_indirection), - current_path=d2h_copier(device_slices.current_path), - log_probs=log_probs_host, - ) - - def _builder() -> BeamHistory | None: - if not need_history.item(): - return None - - if host_snapshot is not None: - cache_indirection = host_snapshot.cache_indirection - current_path = host_snapshot.current_path - log_probs_host = host_snapshot.log_probs - else: - # Predictor-miss fallback: synchronous .cpu() on the main stream. - cache_indirection = device_slices.cache_indirection.cpu() - current_path = device_slices.current_path.cpu() - log_probs_host = None - if device_slices.log_probs is not None: - log_probs_host = _BeamHistoryLogProbsSlices( - sampled_log_probs=device_slices.log_probs.sampled_log_probs.cpu(), - sampled_logprobs_indices=( - device_slices.log_probs.sampled_logprobs_indices.cpu() - ), - cum_logprobs=device_slices.log_probs.cum_logprobs.cpu(), - ) - - new_path = _gather_beam_path( - current_path=current_path, cache_indirection=cache_indirection - ) - new_logprobs: torch.Tensor | None = None - new_logprobs_indices: torch.Tensor | None = None - cum_logprobs_out: torch.Tensor | None = None - if log_probs_host is not None: - new_logprobs, new_logprobs_indices, cum_logprobs_out = _postprocess_beam_logprobs( - request, - cache_indirection=cache_indirection, - log_probs_host=log_probs_host, - ) - - return BeamHistory( - tokens=new_path, - logprobs=new_logprobs, - logprobs_indices=new_logprobs_indices, - cum_logprobs=cum_logprobs_out, - ) - - return _builder + assert cba_group is not None + return _prepare_beam_history_cba(request, cba_group=cba_group) def prepare_beam_histories( self, diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index e1cbc7725577..13bb1898b9cf 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -87,13 +87,7 @@ from ..llm_request import LlmRequest, LlmRequestState, get_draft_token_length from ..resource_manager import ResourceManager, ResourceManagerType from ..scheduler import ScheduledRequests -from .beam_search import ( - BeamHistoryBuilder, - BeamSearchHandler, - _finalize_beam, - _prepare_beam_search, - _request_uses_cba, -) +from .beam_search import BeamHistoryBuilder, BeamSearchHandler, _finalize_beam, _prepare_beam_search from .finish_reasons import FinishReasonsHandler from .logprobs import ( LogProbsState, @@ -2075,8 +2069,7 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: assert beam_search_store is not None # Allocate the CBA tensors on the first beam-search request: every # early_stopping mode runs on that path. - if any(_request_uses_cba(request) for request in new_requests): - beam_search_store.ensure_cba() + beam_search_store.ensure_cba() _prepare_beam_search( beam_search_store, self.store.log_probs_store, diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py index 6b18cbccd1ad..a14a970436bb 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py @@ -35,7 +35,6 @@ BeamSearchMetadata, BeamSearchStore, CBAState, - beam_search_sampling_batch, beam_search_sampling_batch_cba, ) @@ -86,7 +85,6 @@ "CBAState", "Fusions", "StrategyMetadata", - "beam_search_sampling_batch", "beam_search_sampling_batch_cba", "get_rejected_indices", "greedy_search_sampling_batch", @@ -340,7 +338,7 @@ def sample( ): row_stride = cast(BeamSearch, strategy).row_stride assert group_metadata is not None and isinstance(group_metadata, BeamSearchMetadata), ( - "BeamSearchMetadata is required for beam_search_sampling_batch" + "BeamSearchMetadata is required for beam search" ) # Every early_stopping mode goes through the candidate-beams-array # path: TRUE differs only in the done verdict (pool full, without diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 27514d69e5a9..ba66cdb2cc86 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -41,8 +41,7 @@ CBAGroupHost, _finalize_beam, _gather_beam_path, _prepare_beam_history_cba) from tensorrt_llm._torch.pyexecutor.sampler.sampler_strategy import ( BEAM_SEARCH_PAD_TOKEN, BeamSearch, BeamSearchEarlyStop, BeamSearchMetadata, - CBAState, _StrategyImpls, beam_search_sampling_batch, - beam_search_sampling_batch_cba) + CBAState, _StrategyImpls, beam_search_sampling_batch_cba) from tensorrt_llm.bindings.executor import FinishReason from tensorrt_llm.bindings.internal.batch_manager import \ LlmRequest as CppLlmRequest @@ -783,208 +782,6 @@ class GeneralTestParams: vocab_size = 100 -@_kernel_test -def test_beam_search_sampling_batch_basic(): - """Test basic beam search sampling functionality.""" - - test_params = GeneralTestParams() - batch_size = test_params.batch_size - beam_width = test_params.beam_width - vocab_size = test_params.vocab_size - max_batch_size = test_params.max_batch_size - seq_len = test_params.seq_len - temperature = 1.0 - - # Create logits: [batch_size * beam_width, vocab_size] - torch.manual_seed(42) - logits = torch.randn((batch_size * beam_width, vocab_size), - dtype=torch.float32) - for entry in range(batch_size * beam_width): - assert (logits[entry] != logits[entry, 0]).sum( - ) > 0, "Logits of a sequence must not only contain the same value. Otherwise change the seed." - - # create a randomly filled cache indirection - cache_indirection = torch.randint( - 0, - beam_width, - ( - max_batch_size, beam_width, seq_len + 1 - ), # +1 as we should not be calling sampling, when seq_len is already at the max - dtype=torch.int32) - assert cache_indirection.sum( - ) > 0, "Cache indirection must not only contain zeros. Otherwise change the seed." - # create a result tensor for the cache indirection that will be updated by the beam search sampling - cache_indirection_result = cache_indirection.clone() - # Fill this buffer with invalid values - cache_indirection_buffer = torch.full( - (max_batch_size, beam_width, seq_len + 1), -1, dtype=torch.int32) - - # create a zero filled cumulative log probs - cum_log_probs = torch.zeros((max_batch_size, beam_width), - dtype=torch.float32) - # create a result tensor for the cumulative log probs that will be updated by the beam search sampling - cum_log_probs_result = cum_log_probs.clone() - - # add an offset, so that seq slots is not just the first few entries - seq_slots = torch.arange( - batch_size, dtype=torch.int64) + (max_batch_size - batch_size) // 2 - seq_lens = torch.full((batch_size, ), seq_len, dtype=torch.int32) - - # we will finish the last beam of the first request (Note: This also enforces a beam swap in the first request) - finished_beams = torch.zeros((max_batch_size, beam_width), - dtype=torch.int32) - finished_beams[seq_slots[0], beam_width - 1] = FinishReason.STOP_WORDS.value - finished_beams_result = finished_beams.clone() - - new_log_probs = torch.zeros((max_batch_size, beam_width), - dtype=torch.float32) - predecessor_beams = torch.zeros((max_batch_size, beam_width), - dtype=torch.int32) - - new_log_probs_result = new_log_probs.clone() - predecessor_beams_result = predecessor_beams.clone() - - # Pre-computed per-step constants (matches production cache contract). - seq_offsets = torch.arange(max_batch_size, dtype=torch.int64) * beam_width - beam_idx_arange = torch.arange(beam_width, dtype=torch.int32) - - # Create BeamSearchMetadata - beam_search_args = BeamSearchMetadata( - cache_indirection=cache_indirection_result, - cache_indirection_buffer=cache_indirection_buffer, - cum_log_probs=cum_log_probs_result, - seq_slots=seq_slots, - seq_lens=seq_lens, - finished_beams=finished_beams_result, - new_log_probs=new_log_probs_result, - predecessor_beams=predecessor_beams_result, - seq_offsets=seq_offsets, - beam_idx_arange=beam_idx_arange, - beam_gen_lengths=torch.zeros((max_batch_size, beam_width), - dtype=torch.int32), - ) - - # Run beam search sampling. No assert_no_cuda_sync guard here: the guard - # trips on the caching allocator's first allocation for a new shape, not on - # the step itself -- ops that report as synchronizing without a warmed - # allocator (softmax, the scalar-float temperature division, topk) all pass - # the guard once the allocator has served those shapes. Tests that do want - # the no-sync contract go through run_test_with_warmup, which pre-allocates - # for exactly this reason; this one exercises the op's arithmetic instead. - next_tokens, softmax = beam_search_sampling_batch( - logits=logits, - beam_width_in=beam_width, - beam_width_out=beam_width, - beam_search_args=beam_search_args, - temperature=temperature, - return_probs=True, - ) - - # Validate output shapes - assert softmax is not None - expected_tokens_shape = (batch_size, beam_width) - assert next_tokens.shape == expected_tokens_shape, ( - f"Expected shape {expected_tokens_shape}, got {next_tokens.shape}") - expected_softmax_shape = (batch_size, beam_width, vocab_size) - assert softmax.shape == expected_softmax_shape, ( - f"Expected shape {expected_softmax_shape}, got {softmax.shape}") - - # Validate tokens are within vocab range - assert torch.all(next_tokens[1:] >= 0), "Tokens out of vocab range" - # First request has finished beams. Some beams may have BEAM_SEARCH_PAD_TOKEN (-1) as a token - assert torch.all( - next_tokens[0] >= BEAM_SEARCH_PAD_TOKEN), "Tokens out of vocab range" - assert torch.all(next_tokens < vocab_size), "Tokens out of vocab range" - - # Validate softmax probabilities sum to 1 - torch.testing.assert_close(softmax.sum(dim=-1), - torch.ones(batch_size, beam_width)) - - # Setup our expected values: - - # get the top tokens and beams for each request - logprobs = torch.log_softmax(logits, dim=-1) - # adjust our expected log probs accordingly - logprobs[beam_width - 1] = float('-inf') - end_id = vocab_size - 1 - logprobs[beam_width - 1, end_id] = 0 - logprobs = logprobs.view(batch_size, beam_width * vocab_size) - - _, top_indices = torch.topk(logprobs, k=beam_width, dim=-1, sorted=True) - - top_tokens = top_indices % vocab_size - top_beams = top_indices // vocab_size - - torch.cuda.synchronize() - - # Validate cache indirection was updated - for req_idx, seq_slot in enumerate(seq_slots): - for beam_idx in range(beam_width): - ideal_beam = top_beams[req_idx][beam_idx] - torch.testing.assert_close( - cache_indirection_result[seq_slot, beam_idx, :seq_len], - cache_indirection[seq_slot, ideal_beam, :seq_len]) - assert cache_indirection_result[ - seq_slot, beam_idx, - seq_len] == beam_idx, "The last slot of the cache indirection should equal beam_idx" - # Validate cache indirection buffer was updated - for req_idx, seq_slot in enumerate(seq_slots): - for beam_idx in range(beam_width): - torch.testing.assert_close( - cache_indirection_buffer[seq_slot, beam_idx, :], - cache_indirection[seq_slot, beam_idx, :]) - # Validate cumulative log probs were updated - for req_idx, seq_slot in enumerate(seq_slots): - for beam_idx in range(beam_width): - predecessor_beam = top_beams[req_idx][beam_idx] - old_scores = cum_log_probs[seq_slot, predecessor_beam] - new_scores = cum_log_probs_result[seq_slot, beam_idx] - if finished_beams_result[ - seq_slot, beam_idx] != FinishReason.NOT_FINISHED.value: - # finished beams do not calculate the logprobs but set the end tokens logprobs to 0 - torch.testing.assert_close(new_scores, old_scores + 0.0) - else: - torch.testing.assert_close( - new_scores, old_scores + torch.log_softmax( - logits[req_idx * beam_width + predecessor_beam], - dim=-1)[top_tokens[req_idx][beam_idx]]) - # Validate finished beams follow the predecessor-beam swap. The first - # request has a finished beam at column beam_width-1 (set above), so this - # exercises a non-zero finish reason being gathered to its new column. - for req_idx, seq_slot in enumerate(seq_slots): - for beam_idx in range(beam_width): - predecessor_beam = top_beams[req_idx][beam_idx] - torch.testing.assert_close( - finished_beams_result[seq_slot, beam_idx], - finished_beams[seq_slot, predecessor_beam]) - # The finished marker must actually appear post-swap at the column(s) whose - # predecessor was the finished beam, and nowhere else, so a swap that - # dropped or misplaced it would fail. - for req_idx, seq_slot in enumerate(seq_slots): - for beam_idx in range(beam_width): - predecessor_beam = top_beams[req_idx][beam_idx] - expect_finished = (finished_beams[seq_slot, predecessor_beam] - != FinishReason.NOT_FINISHED.value) - got_finished = (finished_beams_result[seq_slot, beam_idx] - != FinishReason.NOT_FINISHED.value) - assert bool(got_finished) == bool(expect_finished) - # test the new log probs - for req_idx, seq_slot in enumerate(seq_slots): - for beam_idx in range(beam_width): - predecessor_beam = top_beams[req_idx][beam_idx] - torch.testing.assert_close( - cum_log_probs_result[seq_slot, beam_idx] - - new_log_probs_result[seq_slot, beam_idx], - cum_log_probs[seq_slot, predecessor_beam]) - # test the predecessor beams - for req_idx, seq_slot in enumerate(seq_slots): - for beam_idx in range(beam_width): - predecessor_beam = top_beams[req_idx][beam_idx] - torch.testing.assert_close( - predecessor_beams_result[seq_slot, beam_idx], - torch.tensor(predecessor_beam, dtype=torch.int32)) - - _CBA_NEG_INF = float("-inf") @@ -1043,115 +840,6 @@ def _with_cba(meta: BeamSearchMetadata, )) -@_kernel_test -@pytest.mark.parametrize("penalty_as_tensor", [False, True]) -def test_beam_search_sampling_batch_length_penalty(penalty_as_tensor): - """Length penalty must flip the ranking between a short finished beam and a - longer active beam with a better per-token score, while stored - cum_log_probs stay raw (unnormalized).""" - - batch_size = 1 - beam_width = 2 - vocab_size = 4 - max_batch_size = 3 - seq_len = 8 - - seq_slots = torch.arange(batch_size, dtype=torch.int64) + 1 - slot = seq_slots[0] - seq_offsets = torch.arange(max_batch_size, dtype=torch.int64) * beam_width - beam_idx_arange = torch.arange(beam_width, dtype=torch.int32) - - def make_metadata() -> BeamSearchMetadata: - cum_log_probs = torch.zeros((max_batch_size, beam_width), - dtype=torch.float32) - # beam 0: finished after 2 generated tokens, frozen cum logprob -1.0 - # beam 1: active with 4 generated tokens, cum logprob -2.0 - cum_log_probs[slot] = torch.tensor([-1.0, -2.0]) - finished_beams = torch.zeros((max_batch_size, beam_width), - dtype=torch.int32) - finished_beams[slot, 0] = FinishReason.END_ID.value - beam_gen_lengths = torch.zeros((max_batch_size, beam_width), - dtype=torch.int32) - beam_gen_lengths[slot] = torch.tensor([2, 4], dtype=torch.int32) - return BeamSearchMetadata( - cache_indirection=torch.zeros( - (max_batch_size, beam_width, seq_len + 1), dtype=torch.int32), - cache_indirection_buffer=torch.full( - (max_batch_size, beam_width, seq_len + 1), - -1, - dtype=torch.int32), - cum_log_probs=cum_log_probs, - seq_slots=seq_slots, - seq_lens=torch.full((batch_size, ), seq_len, dtype=torch.int32), - finished_beams=finished_beams, - new_log_probs=torch.zeros((max_batch_size, beam_width), - dtype=torch.float32), - predecessor_beams=torch.zeros((max_batch_size, beam_width), - dtype=torch.int32), - seq_offsets=seq_offsets, - beam_idx_arange=beam_idx_arange, - beam_gen_lengths=beam_gen_lengths, - ) - - # The active beam (row 1) strongly prefers token 1: logprob ~ -6e-5, so its - # best candidate has raw cum logprob ~ -2.0 vs the finished beam's -1.0. - logits = torch.zeros((batch_size * beam_width, vocab_size), - dtype=torch.float32) - logits[1, 1] = 10.0 - - def run(length_penalty): - metadata = make_metadata() - next_tokens, _ = beam_search_sampling_batch( - logits=logits, - beam_width_in=beam_width, - beam_width_out=beam_width, - beam_search_args=metadata, - temperature=1.0, - length_penalty=length_penalty, - return_probs=False, - ) - return next_tokens, metadata - - # Without penalty: raw scores -1.0 (finished) > ~-2.0 (active) => finished - # beam ranks first and emits a pad token. - tokens_np, meta_np = run(0.0 if not penalty_as_tensor else None) - assert tokens_np[0, 0] == BEAM_SEARCH_PAD_TOKEN - assert tokens_np[0, 1] == 1 - torch.testing.assert_close(meta_np.cum_log_probs[slot, 0], - torch.tensor(-1.0)) - - # With penalty 1.0: normalized scores -1.0/2 = -0.5 (finished) vs - # ~-2.0/5 = -0.4 (active) => the longer active beam ranks first. - penalty = (torch.ones(batch_size, dtype=torch.float32) - if penalty_as_tensor else 1.0) - tokens_lp, meta_lp = run(penalty) - assert tokens_lp[0, 0] == 1 - assert tokens_lp[0, 1] == BEAM_SEARCH_PAD_TOKEN - - # Stored cum_log_probs must remain raw, only the ranking key is normalized. - torch.testing.assert_close(meta_lp.cum_log_probs[slot, 0], - torch.tensor(-2.0), - atol=1e-3, - rtol=1e-3) - torch.testing.assert_close(meta_lp.cum_log_probs[slot, 1], - torch.tensor(-1.0)) - # Generated lengths: active successor grew to 5, finished stays frozen at 2. - assert meta_lp.beam_gen_lengths[slot].tolist() == [5, 2] - # Finished flag must follow the beam swap. - assert meta_lp.finished_beams[slot].tolist() == [ - FinishReason.NOT_FINISHED.value, FinishReason.END_ID.value - ] - - # Negative penalty penalizes the longer beam further: normalized scores - # -1.0 * 2 = -2.0 (finished) vs ~-2.0 * 5 = ~-10 (active), so the short - # finished beam ranks first, same direction as no penalty but stronger. - neg = (-torch.ones(batch_size, dtype=torch.float32) - if penalty_as_tensor else -1.0) - tokens_neg, _ = run(neg) - assert tokens_neg[0, 0] == BEAM_SEARCH_PAD_TOKEN - assert tokens_neg[0, 1] == 1 - - @_kernel_test @pytest.mark.parametrize("params_as_tensors", [False, True]) @pytest.mark.parametrize( @@ -1236,108 +924,6 @@ def test_beam_topk_flashinfer_parity(): assert torch.equal(fi_i[finite], ref_i[finite]) -@_kernel_test -@_kernel_test -@pytest.mark.parametrize("length_penalty", [0.0, 1.0]) -def test_vbws_width_transition_with_length_penalty(length_penalty: float): - """A VBWS width change must keep per-beam lengths aligned to their beams. - - Widening from 2 to 3 beams reorders and duplicates source beams; the - per-beam generated lengths that length_penalty normalizes by must follow - the same permutation. Beams here have deliberately unequal lengths, so a - per-request (rather than per-beam) length would score them identically. - """ - batch_size = 1 - beam_width_in, beam_width_out = 2, 3 - vocab_size = 8 - max_batch_size = 2 - max_beam_width = 4 - seq_len = 6 - device = torch.device("cuda") - - seq_slots = torch.arange(batch_size, dtype=torch.int64, device=device) - slot = int(seq_slots[0]) - - cum_log_probs = torch.zeros((max_batch_size, max_beam_width), - dtype=torch.float32, - device=device) - # Equal cumulative scores, unequal lengths: with length_penalty > 0 the - # longer beam normalizes to a better score, so ranking must change. - cum_log_probs[slot, :beam_width_in] = torch.tensor([-2.0, -2.0], - device=device) - beam_gen_lengths = torch.zeros((max_batch_size, max_beam_width), - dtype=torch.int32, - device=device) - beam_gen_lengths[slot, :beam_width_in] = torch.tensor([2, 5], - dtype=torch.int32, - device=device) - - metadata = BeamSearchMetadata( - cache_indirection=torch.zeros( - (max_batch_size, max_beam_width, seq_len + 1), - dtype=torch.int32, - device=device), - cache_indirection_buffer=torch.full( - (max_batch_size, max_beam_width, seq_len + 1), - -1, - dtype=torch.int32, - device=device), - cum_log_probs=cum_log_probs, - seq_slots=seq_slots, - seq_lens=torch.full((batch_size, ), - seq_len, - dtype=torch.int32, - device=device), - finished_beams=torch.zeros((max_batch_size, max_beam_width), - dtype=torch.int32, - device=device), - new_log_probs=torch.zeros((max_batch_size, max_beam_width), - dtype=torch.float32, - device=device), - predecessor_beams=torch.zeros((max_batch_size, max_beam_width), - dtype=torch.int32, - device=device), - seq_offsets=torch.zeros((batch_size + 1, ), - dtype=torch.int32, - device=device), - beam_idx_arange=torch.arange(max_beam_width, - dtype=torch.int32, - device=device), - beam_gen_lengths=beam_gen_lengths, - ) - - logits = torch.zeros((batch_size * beam_width_in, vocab_size), - dtype=torch.float32, - device=device) - next_tokens, _ = beam_search_sampling_batch( - logits=logits, - beam_width_in=beam_width_in, - beam_width_out=beam_width_out, - beam_search_args=metadata, - temperature=1.0, - length_penalty=length_penalty, - return_probs=False, - ) - - # Rows are laid out at max_beam_width; only the first beam_width_out - # columns carry this step's beams. - assert next_tokens.shape[1] == max_beam_width - # The widened slots must be initialized, not left stale. - assert torch.isfinite(metadata.cum_log_probs[slot, :beam_width_out]).all() - # Slots beyond the current width stay untouched. - torch.testing.assert_close( - metadata.cum_log_probs[slot, beam_width_out:], - torch.zeros(max_beam_width - beam_width_out, device=device)) - - predecessors = metadata.predecessor_beams[slot, :beam_width_out] - assert int(predecessors.min()) >= 0 - assert int(predecessors.max()) < beam_width_in - if length_penalty > 0.0: - # Equal raw scores, so normalization by the per-beam length decides: - # the longer source beam (index 1, length 5) must win. - assert int(predecessors[0]) == 1 - - def test_beam_search_sampling_batch_diversity_rate(): """diversity_rate wired through beam_search_sampling_batch changes beam selection while stored cum_log_probs stay raw.""" From ebdafe004ab9e4759d99ee92ae3a26563e1b2823 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Wed, 5 Aug 2026 22:45:15 -0700 Subject: [PATCH 20/36] [TRTLLM-13234][chore] Remove state left behind by the pool-free path 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 --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 1 - .../_torch/pyexecutor/sampler/beam_search.py | 100 ++---------------- .../_torch/pyexecutor/sampler/sampler.py | 4 +- .../_torch/sampler/test_beam_search.py | 9 -- 4 files changed, 10 insertions(+), 104 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index cefea423f0bf..e30deae6c01d 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -6501,7 +6501,6 @@ def fail_request(message: str) -> bool: # is admitted here, so seed it to 1: length_penalty normalizes by this # counter, and leaving it at zero would divide by one less than the # true length for the whole request. - beam_search_store.beam_gen_lengths[seq_slot, :beam_width].fill_(1) return True @staticmethod diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py index 6a531bacb753..3491c09438b5 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py @@ -15,8 +15,7 @@ """PyTorch-native beam-search sampling kernels. The candidate-selection, beam-expansion, and candidate-beams-array (CBA) -exhaustive-early-stopping logic for TorchSampler beam search, split out of -``ops.vanilla`` so the regular and CBA paths live in one place. Pure tensor +beam-search logic for TorchSampler, split out of ``ops.vanilla``. Pure tensor functions plus their metadata dataclasses; no dependency on the ``sampling_utils`` interface. """ @@ -32,7 +31,7 @@ from tensorrt_llm.bindings.executor import FinishReason from ..llm_request import LlmRequest, LlmRequestState -from .logprobs import LogProbsStore, convert_logprobs_tensor_to_list, get_logprobs_from_request +from .logprobs import LogProbsStore, convert_logprobs_tensor_to_list from .ops.flashinfer import radix_topk_op from .ops.vanilla import StrategyMetadata from .sampler_common import _get_beam_width_in, _unwrap_singleton, int_tensor @@ -60,8 +59,7 @@ class BeamSearchEarlyStop(IntEnum): """ TRUE = 1 - """HF ``True`` (default): stop once ``beam_width`` finished candidates exist. - Regular (non-CBA) path.""" + """HF ``True`` (default): stop once ``beam_width`` finished candidates exist.""" FALSE = 0 """HF ``False``: CBA path bounding attainability by the beam's current @@ -142,8 +140,8 @@ class _CBAFields: @dataclass(kw_only=True) class CBAState(_CBAFields): - """Candidate-Beams-Array (CBA) state, present only for requests using - every beam-search request, whatever its early_stopping mode; see + """Candidate-Beams-Array (CBA) state, present for every beam-search + request whatever its early_stopping mode; see beam_search_sampling_batch_cba. A per-step view over the persistent ``BeamSearchStore`` (shared CBA tensors @@ -197,17 +195,9 @@ class BeamSearchStore: predecessor_beams: torch.Tensor """[max_num_sequences, max_beam_width] int32, predecessor beam per beam, used for stop word detection.""" - seq_offsets: torch.Tensor - """[max_num_sequences] int64, cached ``arange(max_num_sequences) * - max_beam_width`` used by ``beam_search_sampling_batch_cba`` to flatten - (batch_idx, beam_idx) pairs.""" beam_idx_arange: torch.Tensor """[max_beam_width] int32, cached ``arange(max_beam_width)`` used as the scatter source in the per-step ``cache_indirection.scatter_``.""" - beam_gen_lengths: torch.Tensor - """[max_num_sequences, max_beam_width] int32, number of generated tokens per - beam (frozen once a beam finishes). Only maintained for requests with a - non-zero beam-search length_penalty.""" original_tokens: torch.Tensor """[max_num_sequences, max_beam_width, max_seq_len] int32, uncorrected per-slot tokens, written every beam-search step; read with @@ -242,11 +232,7 @@ def create( predecessor_beams=int_tensor(per_beam), original_tokens=int_tensor(cache_indirection_shape), first_finish_reasons=int_tensor(per_beam), - seq_offsets=( - torch.arange(max_num_sequences, device="cuda", dtype=torch.int64) * max_beam_width - ), beam_idx_arange=torch.arange(max_beam_width, device="cuda", dtype=torch.int32), - beam_gen_lengths=int_tensor(per_beam), prompt_lens=int_tensor((max_num_sequences,)), batch_dones=torch.zeros((max_num_sequences,), device="cuda", dtype=torch.bool), ) @@ -290,29 +276,6 @@ class BeamHistory: """[num_beams] float, cumulative log-prob per beam.""" -@dataclass(kw_only=True, frozen=True) -class _BeamHistoryLogProbsSlices: - """Correlated beam-history log-prob tensors; all three fields are bound together.""" - - sampled_log_probs: torch.Tensor - sampled_logprobs_indices: torch.Tensor - cum_logprobs: torch.Tensor - - -@dataclass(kw_only=True, frozen=True) -class _BeamHistoryTensors: - """Beam-history tensor slices. - - Used to carry both device-side views (before D2H) and host-side - snapshots (after D2H). `log_probs` is bound iff log-probs are - requested. - """ - - cache_indirection: torch.Tensor - current_path: torch.Tensor - log_probs: _BeamHistoryLogProbsSlices | None - - def _gather_beam_path( *, current_path: torch.Tensor, cache_indirection: torch.Tensor ) -> torch.Tensor: @@ -334,13 +297,11 @@ class BeamSearchMetadata(StrategyMetadata): seq_lens: torch.Tensor finished_beams: torch.Tensor predecessor_beams: torch.Tensor - seq_offsets: torch.Tensor beam_idx_arange: torch.Tensor - beam_gen_lengths: torch.Tensor stop_past_tokens: Optional[torch.Tensor] = None """[max_stop_word_length, max_num_sequences, max_beam_width] int32, the finish handler's rolling stop-word window (FinishReasonsHandler store). - Used by both the regular and CBA paths: the beam axis is reordered by the + The beam axis is reordered by the step's predecessor beams so the handler's stop-word matching stays correct across beam swaps. When None (tests without stop words), the reorder is skipped — multi-token stop-word matching would then be unreliable across @@ -449,8 +410,8 @@ def beam_candidate_topk( The diversity term ``diversity_rate * source_beam_index`` spreads selection across source beams (candidates from lower-ranked beams get a boost), and the length term normalizes by ``gen_length**length_penalty``. Finished - beams stay frozen in their slots rather than in a separate pool, so a - finished beam's (single) candidate receives the same ``rate * slot_index`` + beams are harvested into the candidate pool, but until the harvest lands + a finished beam's (single) candidate receives the same ``rate * slot_index`` boost as any other — slot position thus slightly affects how hard a finished beam is to evict. @@ -945,7 +906,7 @@ def beam_search_sampling_batch_cba( if reordered_window is not None and stop_window is not None: stop_window[:, slots, :num_beams] = reordered_window - # --- Beam-slot state updates (same contract as beam_search_sampling_batch). + # --- Beam-slot state updates (same contract as beam_search_sampling_batch_cba). args.predecessor_beams[slots, :num_beams] = slot_pred cache_indirection = args.cache_indirection[slots, :num_beams] cache_indirection_buffer = args.cache_indirection_buffer[slots, :beam_width_in] @@ -1031,7 +992,6 @@ def _prepare_beam_search( 0, seq_slots_long, FinishReason.NOT_FINISHED.value ) beam_search_store.original_tokens.index_fill_(0, seq_slots_long, 0) - beam_search_store.beam_gen_lengths.index_fill_(0, seq_slots_long, 0) beam_search_store.prompt_lens.index_copy_(0, seq_slots_long, prompt_lens_cuda) beam_search_store.batch_dones.index_fill_(0, seq_slots_long, False) # The CBA tensors only exist once a beam-search request has @@ -1155,48 +1115,6 @@ def _builder() -> BeamHistory | None: return _builder -def _postprocess_beam_logprobs( - request: LlmRequest, - *, - cache_indirection: torch.Tensor, - log_probs_host: _BeamHistoryLogProbsSlices, -) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Reorder per-step beam logprobs along the cache-indirection axis. - - Concatenates the freshly-sampled per-step entries onto the - request's existing host-side logprobs buffer and gathers each - beam's history through `cache_indirection`. Returns the gathered - (logprobs, logprobs_indices, cum_logprobs) triple. - """ - current_logprobs, current_logprobs_indices = get_logprobs_from_request( - request, preallocate_extra_steps=1 - ) - # concatenate the newly generated logprobs and newly - # generated tokens to the current logprobs and logprobs indices - current_logprobs[:, -1, :].copy_(log_probs_host.sampled_log_probs) - current_logprobs_indices[:, -1, :].copy_(log_probs_host.sampled_logprobs_indices) - - # Gather the correct logprobs for each beam. - new_logprobs = torch.zeros_like(current_logprobs) - new_logprobs_indices = torch.zeros_like(current_logprobs_indices) - cache_indirection_for_logprobs = cache_indirection.unsqueeze(-1).expand( - -1, -1, current_logprobs.shape[2] - ) - torch.gather( - input=current_logprobs, - dim=0, - index=cache_indirection_for_logprobs, - out=new_logprobs, - ) - torch.gather( - input=current_logprobs_indices, - dim=0, - index=cache_indirection_for_logprobs, - out=new_logprobs_indices, - ) - return new_logprobs, new_logprobs_indices, log_probs_host.cum_logprobs - - def _finalize_beam( request: LlmRequest, beam_history: BeamHistory, diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 13bb1898b9cf..2cb0b63131aa 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -2258,9 +2258,7 @@ def _add_metadata_to_grouped_requests( seq_lens=group_seq_lens_cuda, finished_beams=beam_search_store.first_finish_reasons, predecessor_beams=beam_search_store.predecessor_beams, - seq_offsets=beam_search_store.seq_offsets, beam_idx_arange=beam_search_store.beam_idx_arange, - beam_gen_lengths=beam_search_store.beam_gen_lengths, stop_past_tokens=self._finish_reasons_handler.store.past_tokens_cuda, # None unless a beam-search request has been # admitted; the CBA tensors are not allocated before that. @@ -3632,7 +3630,7 @@ def _gather_src_dst_indices( # Because req_num_generated_tokens may differ from the number of sampled tokens in # beam search, the sampled rank computation would entail extra complexity to resolve # the relationships between incoming and outgoing beams. For the sampled logprobs, this - # matching happens in beam_search_sampling_batch() which updates + # matching happens in beam_search_sampling_batch_cba() which updates # log_probs_store.sampled_log_probs. Therefore, neither sampled ranks nor sampled logprobs # are handled here. if logprobs_reqs_indices_n_beam: diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index ba66cdb2cc86..2f9932ddc30b 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -955,11 +955,7 @@ def make_metadata() -> BeamSearchMetadata: dtype=torch.float32), predecessor_beams=torch.zeros((max_batch_size, beam_width), dtype=torch.int32), - seq_offsets=torch.arange(max_batch_size, dtype=torch.int64) * - beam_width, beam_idx_arange=torch.arange(beam_width, dtype=torch.int32), - beam_gen_lengths=torch.zeros((max_batch_size, beam_width), - dtype=torch.int32), ) logits = torch.full((batch_size * beam_width, vocab_size), -20.0) @@ -1026,9 +1022,7 @@ def _make_cba_metadata(max_batch, K, attn_len, snap_len, seq_len, prompt_len, seq_lens=torch.full((batch, ), seq_len, dtype=torch.int32), finished_beams=torch.zeros((max_batch, K), dtype=torch.int32), predecessor_beams=torch.zeros((max_batch, K), dtype=torch.int32), - seq_offsets=torch.arange(max_batch, dtype=torch.int64) * K, beam_idx_arange=torch.arange(K, dtype=torch.int32), - beam_gen_lengths=torch.zeros((max_batch, K), dtype=torch.int32), cba=CBAState( end_ids=torch.full((max_batch, ), end_id, dtype=torch.int32), prompt_lens=torch.full((max_batch, ), prompt_len, @@ -1206,9 +1200,6 @@ def run(length_penalty): batch=1) for b in range(K): m.cache_indirection[0, b, :] = b - # beam0 is the short one: it already finished two tokens ago, so its - # frozen length is smaller than beam1's. - m.beam_gen_lengths[0] = torch.tensor([1, gen], dtype=torch.int32) # Raw scores: beam0 ahead of beam1. m.cum_log_probs[0] = torch.tensor([-1.0, -2.0]) From 512a13d282ec6282ebd84905593cf24ab182d6ad Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Thu, 6 Aug 2026 01:13:12 -0700 Subject: [PATCH 21/36] [TRTLLM-13234][feat] Gate speculative beam-history D2H per group 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 --- .../_torch/pyexecutor/sampler/beam_search.py | 50 +++++++++++++++++ .../test_beam_search_speculative_d2h.py | 53 +++++++++++++++++++ 2 files changed, 103 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py index 3491c09438b5..8c7980f3a36e 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py @@ -1368,6 +1368,26 @@ def prepare_beam_histories( forward it to _record_sampler_event so SamplerEvent.synchronize awaits the side stream before any builder is invoked. """ + # The snapshot is only ever read when a request finishes on this step, + # so in speculative mode skip it entirely on steps where no request can + # finish. The predictor is per request but the copy is per group, so + # the group is skipped only when every request is predicted + # non-terminal; one possible finisher pays for the whole group, which + # is the same copy it would have cost anyway. + issue_copies = not self._use_speculative_d2h or any( + self.predict_is_likely_finishing( + req, + num_generated_tokens=req.max_beam_num_tokens + 1 - req.py_prompt_len, + num_tokens=req.max_beam_num_tokens + 1, + ) + for req in requests + ) + if not issue_copies: + # No snapshot. A builder falls back to a synchronous read if the + # step turns out to have finished a request after all -- see + # _deferred_cba_group. + return [self._speculative_builder(req) for req in requests], None + # Single `with` for both modes; nullcontext yields None. copier_ctx: AbstractContextManager["_SideStreamCopier | None"] = ( self._make_side_stream_copier() if self._use_speculative_d2h else nullcontext() @@ -1388,3 +1408,33 @@ def prepare_beam_histories( ] side_stream_event = copier.event if copier is not None else None return builders, side_stream_event + + def _speculative_builder(self, request: LlmRequest) -> BeamHistoryBuilder: + """Builder for a step whose snapshot was skipped. + + Takes the snapshot synchronously when invoked. `should_stop` inside it + still decides whether a history is produced, so a prediction that held + costs one blocking copy of this request's rows and returns None; a miss + costs the same copy and produces the history. Both stall the step, + which is what the predictor trades against -- it is only worth enabling + when most steps finalize nothing, in which case no builder for a + skipped step is ever invoked in the first place. + + NB: deliberately not gated on a host-side finish check. The mirror the + finish handler keeps (`_prev_first_finish_reasons`) lags by one step, + so a request finishing on this step -- a stop word, say -- would read + as unfinished and its history would be dropped. + """ + + def _builder() -> BeamHistory | None: + store = self._store + assert store is not None + cba_group = self.prepare_cba_group_host( + [request], store.first_finish_reasons, self._copy_to_host + ) + if cba_group is None: + return None + inner = _prepare_beam_history_cba(request, cba_group=cba_group) + return inner() if inner is not None else None + + return _builder diff --git a/tests/unittest/_torch/sampler/test_beam_search_speculative_d2h.py b/tests/unittest/_torch/sampler/test_beam_search_speculative_d2h.py index 1d899c87f7d8..3a5004f30aee 100644 --- a/tests/unittest/_torch/sampler/test_beam_search_speculative_d2h.py +++ b/tests/unittest/_torch/sampler/test_beam_search_speculative_d2h.py @@ -336,6 +336,59 @@ def test_speculative_d2h_predictor_always_hit( _assert_outputs_equal(out_on, out_off) +@pytest.mark.threadleak(enabled=False) +def test_speculative_d2h_skips_only_when_no_request_can_finish( + fixed_params: dict[str, Any], + input_prompts: list[list[int]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The skip decision is per group, not per request. + + The snapshot is one batched copy for the whole group, so it cannot be + skipped for some requests and taken for others. A step is skipped only when + every request is predicted non-terminal; a single possible finisher makes + the group pay for the copy it would have cost anyway. + + Pin the predictor to answer True for exactly one request and False for the + rest. The mixed verdict must land on "copy", and the outputs must still + match the feature-off run -- a per-request reading of the same verdict + would drop the histories of the requests that answered False. + """ + seen_verdicts: list[bool] = [] + first_slot: dict[str, int | None] = {"slot": None} + + def _one_finisher(self, request, *, num_generated_tokens, num_tokens): + # Latch onto whichever request is seen first and only ever claim that + # one might finish. + if first_slot["slot"] is None: + first_slot["slot"] = request.py_seq_slot + verdict = request.py_seq_slot == first_slot["slot"] + seen_verdicts.append(verdict) + return verdict + + with monkeypatch.context() as mp_off: + out_off = _run_with_env( + fixed_params, input_prompts, mp_off, speculative=False, stop_token_ids=None + ) + + with monkeypatch.context() as mp_on: + out_on = _run_with_env( + fixed_params, + input_prompts, + mp_on, + speculative=True, + stop_token_ids=None, + predictor_override=_one_finisher, + ) + + assert seen_verdicts, "predictor patch was never invoked; the speculative path did not run" + assert any(seen_verdicts) and not all(seen_verdicts), ( + "the test needs a step where the verdicts disagree, so that the group " + f"decision is observable; got {set(seen_verdicts)}" + ) + _assert_outputs_equal(out_on, out_off) + + # --------------------------------------------------------------------------- # Validator: speculative path must be rejected when sampler_force_async_worker # is also set, since the speculative path bypasses the async D2H worker. From f8d55892731decdf66f34bb2eacf2316510553a3 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Thu, 6 Aug 2026 01:29:20 -0700 Subject: [PATCH 22/36] [TRTLLM-13234][fix] Address beam-search review findings 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 - 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 --- docs/source/features/sampling.md | 7 +- tensorrt_llm/_torch/pyexecutor/llm_request.py | 10 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 25 ++- .../_torch/pyexecutor/sampler/beam_search.py | 89 +++++----- .../pyexecutor/sampler/sampler_strategy.py | 7 +- .../_torch/sampler/test_beam_search.py | 156 +++++++++++++----- 6 files changed, 198 insertions(+), 96 deletions(-) diff --git a/docs/source/features/sampling.md b/docs/source/features/sampling.md index 9a20e1de47f8..277aa64eeddc 100644 --- a/docs/source/features/sampling.md +++ b/docs/source/features/sampling.md @@ -301,8 +301,11 @@ Beam search rejects the following combinations, raising an error at admission: - **A decreasing `beam_width_array`.** Only non-decreasing schedules are supported; the semantics of narrowing mid-decode are not defined. - **A `best_of` other than `max_beam_width`.** Every request in an engine runs at the same - beam width. Note that mixing widths would fail at forward time rather than per request, so - the check happens on admission instead. + beam width, which admission enforces so that a mismatch is reported against the offending + request. Mixing widths is a forward-time failure that aborts the whole batch: note that + admission compares `best_of` against `max_beam_width` only, so requests whose + `beam_width_array` puts them at different per-iteration widths in the same step still + reach that failure. The following example demonstrates beam search with a beam width of 4, returning the top 3 sequences: diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 04c77de3f99b..9d9a744e1e94 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -885,12 +885,16 @@ def get_beam_width_by_iter(self, for_next_iteration: bool = False) -> int: which the mismatch starves the request in the micro-batch scheduler and decoding hangs. test_vbws_cpp_formula_matches_past_array_end pins the agreement. + + An empty array (``[]`` or ``[[]]``) falls through to the base + implementation rather than indexing it, matching the two emptiness + guards the C++ side checks before reading element 0. """ beam_width_array = self.sampling_config.beam_width_array - if beam_width_array is not None: - if beam_width_array and isinstance(beam_width_array[0], - (list, tuple)): + if beam_width_array: + if isinstance(beam_width_array[0], (list, tuple)): beam_width_array = beam_width_array[0] + if beam_width_array: iteration = self.decoding_iter + (1 if for_next_iteration else 0) index = max(min(iteration, len(beam_width_array)) - 1, 0) return int(beam_width_array[index]) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index e30deae6c01d..ed06c7449733 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -5010,6 +5010,22 @@ def _validate_request(self, request: LlmRequest): "only non-decreasing arrays are supported for " "Variable-Beam-Width-Search.") + # length_penalty is an exponent on the generated length, and the + # ranking assumes it only ever shrinks the magnitude of a negative + # cum_log_prob. A negative exponent inverts that: dividing by + # length**negative multiplies instead, so longer beams score + # higher and the beam order is reversed. + length_penalty = sampling_config.length_penalty + if length_penalty is not None: + if isinstance(length_penalty, (list, tuple)): + invalid = [p for p in length_penalty if p < 0] + else: + invalid = [length_penalty] if length_penalty < 0 else [] + if invalid: + raise ValueError( + f"length_penalty {invalid} is negative; only " + "non-negative values are supported.") + # Beam search keeps a pool of finished candidates, which the # context server can already populate on its first step. That pool # is not part of the disaggregated handoff (ContextPhaseParams @@ -6496,11 +6512,10 @@ def fail_request(message: str) -> bool: dtype=cum_log_probs.dtype) cum_log_probs[seq_slot, :beam_width].copy_(values) - # The handoff carries one token already produced upstream. The - # per-beam generated-length counter is reset to zero when the request - # is admitted here, so seed it to 1: length_penalty normalizes by this - # counter, and leaving it at zero would divide by one less than the - # true length for the whole request. + # The handoff carries one token already produced upstream, seeded above + # at index prompt_len. No generated-length counter needs seeding to + # match: length_penalty normalizes by seq_lens - prompt_lens, which + # counts that token because it occupies the first generated slot. return True @staticmethod diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py index 8c7980f3a36e..faf4bb0abfa8 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py @@ -27,6 +27,7 @@ import torch +from tensorrt_llm._torch.flashinfer_utils import IS_FLASHINFER_AVAILABLE from tensorrt_llm._utils import nvtx_range, prefer_pinned from tensorrt_llm.bindings.executor import FinishReason @@ -93,8 +94,14 @@ def _beam_topk(values: torch.Tensor, k: int) -> tuple[torch.Tensor, torch.Tensor large rows (e.g. vocab-sized logits), while torch.topk wins on small rows where the radix kernel's fixed per-call cost dominates. The 10k crossover follows flashinfer's own guidance (``flashinfer.top_k`` docstring). + + Beam search is not a flashinfer-gated feature, so fall back to torch.topk + when flashinfer is missing: real vocabularies are always past the crossover, + which would otherwise make the whole default beam path require it. The two + kernels agree on values and on descending order; only the index order among + equal values may differ, which torch.topk leaves unspecified anyway. """ - if values.size(-1) > 10000: + if IS_FLASHINFER_AVAILABLE and values.size(-1) > 10000: return radix_topk_op(values, k) return torch.topk(values, k=k, dim=-1, sorted=True) @@ -607,6 +614,21 @@ def _take(src: torch.Tensor) -> torch.Tensor: step_log_probs = slot_cum - cum_log_probs[slots, :beam_width_in].gather(1, slot_pred.long()) # --- CBA insertion: normalized scores of the new end-token candidates. + # beam_candidate_topk applies the diversity term for *ranking* only and + # returns raw cumulative log-probs, so the CBA insertion score below and + # best_attainable further down are both diversity-free. This matches HF, + # where diversity is a logits processor (HammingDiversityLogitsProcessor) + # applied before log_softmax and so never reaches the accumulated score, + # and vLLM, which ranks purely by cum_logprob / seq_len**length_penalty. + # + # The C++ kernels here diverge from all three: they fold + # `diversityRate * source_beam_index` into pLocalLogProbs + # (beamSearchKernels.cu), which then flows into normedScoresCBA and + # bestAttainableScore (beamSearchKernelsTemplate.h), so with diversity_rate + # set they order the pool and reach the done verdict differently. Do not + # "fix" this by matching them: a cumulative_logprob carrying + # `diversity_rate * beam_index` is no longer a log-probability, and the + # offset depends on which slot the beam happened to occupy (TRTLLM-14792). new_normed = cand_cum / cand_len.to(cand_cum.dtype).pow(exponent) eligible = is_end & (cand_rank < caps) new_normed = new_normed.masked_fill(~eligible, neg_inf) @@ -763,7 +785,24 @@ def beam_search_sampling_batch_cba( marking every beam slot finished, which drives the regular stop machinery. - Requires the CBA fields of ``BeamSearchMetadata`` to be set. + Execution contract: + + - **Mutates ``beam_search_args`` in place** and returns only the sampled + tokens (and probabilities). The updated beam state -- ``cum_log_probs``, + ``new_log_probs``, ``finished_beams``, ``predecessor_beams``, + ``cache_indirection``, ``original_tokens``/``original_log_probs``, and + the CBA fields (``cba_tokens``, ``cba_cum_log_probs``, + ``cba_normed_scores``, ``cba_lengths``, ``batch_dones``) -- is read back + from the caller's tensors, not from the return value. Only the rows + selected by ``seq_slots`` are written. + - Requires the CBA fields of ``BeamSearchMetadata`` to be set. + - Calls ``mark_dynamic`` on the caller's tensors so a single compiled + graph serves every batch size. The caller must not rely on those + tensors keeping static shapes afterwards. + - Dispatches between a ``torch.compile``d ``_cba_step_math`` and an eager + fallback. Both compute the same values; the split exists because the + scatter-style write-back does not survive fullgraph tracing, so it runs + outside the compiled region (see the write-back comments below). """ args = beam_search_args cba = args.cba @@ -1317,42 +1356,6 @@ def prepare_cba_group_host( ), ) - def _prepare_beam_history( - self, - request: LlmRequest, - *, - finish_reasons: torch.Tensor, - d2h_copier: Callable[[torch.Tensor], torch.Tensor], - cba_group: Optional[CBAGroupHost] = None, - ) -> BeamHistoryBuilder | None: - """Correct the stored tokens for each beam and return it as a BeamHistory object. - - Beam Search sampling only adds new tokens to the beam. - However during beam search, a beam may change its previously sampled tokens. - This function corrects the stored tokens for each beam to match the expected tokens. - If logprobs are requested, the function also corrects the stored logprobs for each beam. - The function returns a BeamHistory object that contains the corrected tokens and logprobs for each beam. - - D2H copies are issued through `d2h_copier`. When - `_use_speculative_beam_history_d2h` is set, a host-side predictor - decides per step whether to stage copies via `d2h_copier`; - predictor misses fall back to a synchronous `.cpu()` inside - `_builder`. Otherwise, copies are issued unconditionally. - - Note: To defer the decision whether or not to skip BeamHistory construction until update_requests(), only - a builder (BeamHistoryBuilder) is returned here. The builder contains host tensors which are - being populated asynchronously. Hence, it can only be invoked after async D2H copies have completed, - e.g., after awaiting state.sampler_event in update_requests. - - arguments: - request: The request to create the beam history for - finish_reasons: The first finish reason encountered for each beam of the request. - Shape: (max_tokens, max_beam_width) - d2h_copier: Callable performing the D2H copy. - """ - assert cba_group is not None - return _prepare_beam_history_cba(request, cba_group=cba_group) - def prepare_beam_histories( self, requests: list[LlmRequest], @@ -1397,15 +1400,7 @@ def prepare_beam_histories( copier.stage_copy_to_host if copier is not None else self._copy_to_host ) cba_group = self.prepare_cba_group_host(requests, finish_reasons, d2h_copier) - builders = [ - self._prepare_beam_history( - req, - finish_reasons=finish_reasons[req.py_seq_slot], - d2h_copier=d2h_copier, - cba_group=cba_group, - ) - for req in requests - ] + builders = [_prepare_beam_history_cba(req, cba_group=cba_group) for req in requests] side_stream_event = copier.event if copier is not None else None return builders, side_stream_event diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py index a14a970436bb..fc04afbb51e4 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py @@ -38,8 +38,11 @@ beam_search_sampling_batch_cba, ) -# These op wrappers are safe to import without flashinfer installed; they are -# only called on the flashinfer sampler / speculative-worker paths. +# These op wrappers are safe to import without flashinfer installed; each one +# resolves the flashinfer symbol only when called. Most are reached only on the +# flashinfer sampler / speculative-worker paths, but radix_topk_op also backs +# beam search's wide-row top-k, which guards on IS_FLASHINFER_AVAILABLE and +# falls back to torch.topk (see beam_search._beam_topk). from tensorrt_llm._torch.pyexecutor.sampler.ops.flashinfer import ( sampling_from_probs_op, sanitize_top_k, diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 2f9932ddc30b..17976369d77b 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -17,6 +17,7 @@ import gc import os import pathlib as _pl +import types from contextlib import contextmanager, nullcontext from copy import deepcopy from typing import Any, Callable, Generator, cast @@ -34,6 +35,7 @@ from tensorrt_llm._torch.pyexecutor.llm_request import (LlmRequest, LlmRequestState, SamplingConfig) +from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor from tensorrt_llm._torch.pyexecutor.sampler import (BeamHistory, SampleStateTorch, TorchSampler) @@ -670,15 +672,6 @@ def recording_update_requests(self: TorchSampler, state, *args, **kwargs): monkeypatch.setattr(TorchSampler, "update_requests", recording_update_requests) - # Each parametrization builds a fresh engine in this same process, and the - # CBA step is torch.compile'd with fullgraph=True. Dynamo counts recompiles - # per code object across the whole process, so by the last case the default - # limit is exhausted and compilation fails hard rather than falling back. - # Raise it for the duration of the test; the limit is a guard against - # runaway recompilation, not a correctness property. - monkeypatch.setattr(torch._dynamo.config, "recompile_limit", - max(64, torch._dynamo.config.recompile_limit)) - gc.collect(2) # force destruction of any other LLM instances with _single_process_context(): llm = LLM( @@ -751,6 +744,24 @@ def recording_update_requests(self: TorchSampler, state, *args, **kwargs): ########################################################################### +@pytest.fixture(autouse=True) +def _raise_dynamo_recompile_limit(monkeypatch): + """Keep the whole module clear of Dynamo's per-code-object recompile cap. + + Nearly every test here builds a fresh engine or a fresh set of CBA tensors + in this same process, and the CBA step is compiled with fullgraph=True. + Dynamo counts recompiles per code object across the entire process, so the + default cap is exhausted partway through the file and every later case + fails to compile -- a hard failure under fullgraph, not a fallback. Each + of those tests passes when run alone. + + The cap guards against runaway recompilation; it is not a correctness + property, so raising it for the module is safe. + """ + monkeypatch.setattr(torch._dynamo.config, "recompile_limit", + max(256, torch._dynamo.config.recompile_limit)) + + def _kernel_test(fn: Callable[..., Any]) -> Callable[..., Any]: """Mark a beam-search kernel unit test as CUDA-only and run its body with the default device set to CUDA, so the tensors it builds land on GPU (the @@ -924,6 +935,7 @@ def test_beam_topk_flashinfer_parity(): assert torch.equal(fi_i[finite], ref_i[finite]) +@_kernel_test def test_beam_search_sampling_batch_diversity_rate(): """diversity_rate wired through beam_search_sampling_batch changes beam selection while stored cum_log_probs stay raw.""" @@ -1007,17 +1019,54 @@ def run(diversity_rate): rtol=1e-3) +# Poison value for buffers the op is expected to overwrite wherever it writes +# at all. Zero is a plausible real value for most of them (a token id, a +# length, a beam index), so a zero-filled buffer cannot distinguish "written +# correctly" from "never written"; these can. +_UNWRITTEN_INT = -7 +_UNWRITTEN_FLOAT = -7.0 + + +def _assert_untouched_rows_intact(meta, before, batch, max_batch): + """Rows outside seq_slots must come back bit-identical. + + The ops index every buffer with seq_slots, so a slot-agnostic write (a + forgotten index, a broadcast over dim 0) still produces correct output for + the slots under test and is invisible unless the unused rows are checked. + """ + if batch >= max_batch: + return + unused = slice(batch, max_batch) + for name, orig in before.items(): + current = _cba_field(meta, name) + assert torch.equal(current[unused], orig[unused]), ( + f"{name} rows [{batch}:{max_batch}] were modified; the op must " + "only write the rows selected by seq_slots") + + +def _cba_field(meta, name): + return getattr(meta.cba, name) if hasattr(meta.cba, name) else getattr( + meta, name) + + +def _snapshot_cba_rows(meta, names): + return {name: _cba_field(meta, name).clone() for name in names} + + def _make_cba_metadata(max_batch, K, attn_len, snap_len, seq_len, prompt_len, end_id, batch): slots = torch.arange(batch, dtype=torch.int64) m = BeamSearchMetadata( - cache_indirection=torch.zeros((max_batch, K, attn_len), - dtype=torch.int32), + cache_indirection=torch.full((max_batch, K, attn_len), + _UNWRITTEN_INT, + dtype=torch.int32), cache_indirection_buffer=torch.full((max_batch, K, attn_len), -1, dtype=torch.int32), cum_log_probs=torch.zeros((max_batch, K), dtype=torch.float32), - new_log_probs=torch.zeros((max_batch, K), dtype=torch.float32), + new_log_probs=torch.full((max_batch, K), + _UNWRITTEN_FLOAT, + dtype=torch.float32), seq_slots=slots, seq_lens=torch.full((batch, ), seq_len, dtype=torch.int32), finished_beams=torch.zeros((max_batch, K), dtype=torch.int32), @@ -1027,22 +1076,29 @@ def _make_cba_metadata(max_batch, K, attn_len, snap_len, seq_len, prompt_len, end_ids=torch.full((max_batch, ), end_id, dtype=torch.int32), prompt_lens=torch.full((max_batch, ), prompt_len, dtype=torch.int32), - original_tokens=torch.zeros((max_batch, K, attn_len), - dtype=torch.int32), + original_tokens=torch.full((max_batch, K, attn_len), + _UNWRITTEN_INT, + dtype=torch.int32), cba_tokens=torch.full((max_batch, K, snap_len), BEAM_SEARCH_PAD_TOKEN, dtype=torch.int32), - cba_cum_log_probs=torch.zeros((max_batch, K), dtype=torch.float32), + cba_cum_log_probs=torch.full((max_batch, K), + _UNWRITTEN_FLOAT, + dtype=torch.float32), cba_normed_scores=torch.full((max_batch, K), _CBA_NEG_INF, dtype=torch.float32), - cba_lengths=torch.zeros((max_batch, K), dtype=torch.int32), + cba_lengths=torch.full((max_batch, K), + _UNWRITTEN_INT, + dtype=torch.int32), batch_dones=torch.zeros((max_batch, ), dtype=torch.bool), cba_caps=torch.full((max_batch, ), K, dtype=torch.int32), - original_log_probs=torch.zeros((max_batch, K, attn_len), - dtype=torch.float32), - cba_log_probs=torch.zeros((max_batch, K, snap_len), - dtype=torch.float32), + original_log_probs=torch.full((max_batch, K, attn_len), + _UNWRITTEN_FLOAT, + dtype=torch.float32), + cba_log_probs=torch.full((max_batch, K, snap_len), + _UNWRITTEN_FLOAT, + dtype=torch.float32), max_seq_len=attn_len, ), ) @@ -1169,7 +1225,6 @@ def test_beam_search_cba_done_bound_by_early_stopping(early_stopping, -1.5])) -@_kernel_test @_kernel_test @pytest.mark.parametrize("penalty_as_tensor", [False, True]) def test_beam_search_cba_length_penalty_orders_pool(penalty_as_tensor): @@ -1210,6 +1265,11 @@ def run(length_penalty): penalty = (torch.full((1, ), length_penalty, dtype=torch.float32) if penalty_as_tensor else length_penalty) + before = _snapshot_cba_rows( + m, ("cache_indirection", "cum_log_probs", "new_log_probs", + "finished_beams", "predecessor_beams", "original_tokens", + "cba_tokens", "cba_cum_log_probs", "cba_normed_scores", + "cba_lengths", "batch_dones")) beam_search_sampling_batch_cba( logits=logits, beam_width_in=K, @@ -1220,6 +1280,7 @@ def run(length_penalty): length_penalty=penalty, return_probs=False, ) + _assert_untouched_rows_intact(m, before, batch=1, max_batch=2) return m # Both runs put the same two hypotheses in the pool; only the ordering key @@ -1261,6 +1322,7 @@ def run(length_penalty): f"normed {normed} != cum {cum} / length {length}") +@_kernel_test def test_beam_search_cba_replace_min(): """A better finished path replaces the worst CBA entry when full.""" K, vocab, end_id = 2, 5, 4 @@ -1474,18 +1536,19 @@ def _vbws_request(beam_width_array: list[int] | None, @pytest.mark.parametrize( "beam_width_array, expected", [ - # Widening: index is (iteration - 1), clamped at both ends. + # Index is (iteration - 1), clamped at both ends. ([2, 3, 4], [2, 2, 3, 4]), - # Narrowing: beams are dropped as decoding proceeds. - ([4, 3, 2], [4, 4, 3, 2]), + ([2, 2, 4], [2, 2, 2, 4]), ], - ids=["widening", "narrowing"], + ids=["widening", "flat_then_widening"], ) def test_vbws_beam_width_by_iter_follows_array(beam_width_array: list[int], expected: list[int]): """get_beam_width_by_iter walks beam_width_array as decoding advances. - Covers both directions: widening adds beams, narrowing drops them. + Only non-decreasing arrays are covered: a narrowing one is rejected at + admission (test_vbws_rejects_decreasing_beam_width_array), so pinning the + width it would walk through would be pinning unreachable behaviour. """ request = _vbws_request(beam_width_array) actual = [] @@ -1504,11 +1567,12 @@ def test_vbws_beam_width_by_iter_follows_array(beam_width_array: list[int], def test_vbws_beam_width_by_iter_clamps_past_array_end(): """Decoding longer than beam_width_array must hold the last width. - Regression test for the C++ formula, which clamps the iteration index with - the global kMaxBeamWidthArrayLength (assuming a padded array) and therefore - reads past the end of the raw user array, returning garbage. The Python - override clamps with the actual array length instead; see - LlmRequest.get_beam_width_by_iter. + The C++ formula used to clamp the iteration index with the global + kMaxBeamWidthArrayLength (assuming a padded array), reading past the end of + the raw user array and returning garbage; llmRequest.cpp now clamps with the + array's own length. This pins the Python override, which clamps the same way + and is kept so callers do not bind to a prebuilt library from before that + fix; see LlmRequest.get_beam_width_by_iter. """ beam_width_array = [2, 3, 4] request = _vbws_request(beam_width_array) @@ -1571,17 +1635,35 @@ def test_vbws_rejects_decreasing_beam_width_array(beam_width_array: list[int], step. Reject at admission instead of emitting silently stale beams. """ - # Mirrors the check in PyExecutor._validate_request. - def is_accepted(array: list[int]) -> bool: - return not any(b < a for a, b in zip(array, array[1:])) - - assert is_accepted(beam_width_array) == accepted - request = _vbws_request(beam_width_array) # The request itself is still constructible; admission is what rejects it, # and py_beam_width is the array maximum either way. assert request.py_beam_width == max(beam_width_array) + # Drive the real admission check rather than a local re-implementation: + # _validate_request only reads self.max_beam_width, so a stub carrying it + # is enough to reach the beam_width_array branch without standing up an + # executor. A test that mirrored the predicate would keep passing if the + # production check were deleted. + # Everything _validate_request touches besides the beam checks runs after + # them and needs a live engine/sampler, so stub those two out; the beam + # width and beam_width_array branches are reached with the real code. + executor = types.SimpleNamespace( + max_beam_width=request.py_beam_width, + _validate_token_id_range=lambda _request: None, + sampler=types.SimpleNamespace(validate_request=lambda _request: None), + ) + validate = functools.partial( + PyExecutor._validate_request, + executor, # type: ignore[arg-type] # stub: only max_beam_width is read + ) + + if accepted: + validate(request) + else: + with pytest.raises(ValueError, match="decreases"): + validate(request) + def test_beam_strategy_grouping_key_tolerates_trailing_fields(): """The grouping key must not depend on the BeamSearch tuple's arity. From a2ea98142d8b8852853014454c453de1b9049094 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Thu, 6 Aug 2026 03:28:22 -0700 Subject: [PATCH 23/36] [TRTLLM-13234][fix] Keep the earliest per-beam finish reason on the CBA 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 --- .../_torch/pyexecutor/sampler/beam_search.py | 22 +++++++++++++------ 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py index faf4bb0abfa8..e469dda0e720 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py @@ -934,13 +934,21 @@ def beam_search_sampling_batch_cba( cba.cba_tokens[slots, :, :snap_len] = merged_tokens cba.cba_log_probs[slots, :, :snap_len] = merged_lps cba.batch_dones[slots] = done - # NOT_FINISHED == 0, so the verdict maps directly onto the finish reason. - # Flood the full row: the stop criterion reads the first py_beam_width - # (== capacity) entries, which can exceed this step's beam_width_out for - # variable-beam-width requests. - args.finished_beams[slots] = ( - done.view(-1, 1).to(torch.int32) * FinishReason.END_ID.value - ).expand(-1, args.finished_beams.size(1)) + # Publish the done verdict by flooding the full row: the stop criterion + # reads the first py_beam_width (== capacity) entries, which can exceed + # this step's beam_width_out for variable-beam-width requests. + # + # Only write where the verdict is done. This tensor is the finish handler's + # first_finish_reasons, which records the *earliest* reason each beam + # finished by and must never be cleared: a beam that ended on a stop word + # vacates its slot and keeps generating on this path, so an unconditional + # write would erase that STOP_WORDS entry on the next step and the request + # would report the reason it eventually stopped for (LENGTH) instead. + args.finished_beams[slots] = torch.where( + done.view(-1, 1), + torch.full_like(args.finished_beams[slots], FinishReason.END_ID.value), + args.finished_beams[slots], + ) stop_window = args.stop_past_tokens if reordered_window is not None and stop_window is not None: stop_window[:, slots, :num_beams] = reordered_window From 33d2992c5815de4e31274adfea0aa3278530983a Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Thu, 6 Aug 2026 05:33:45 -0700 Subject: [PATCH 24/36] [TRTLLM-13234][fix] Report per-beam finish reasons on the CBA path 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 --- .../_torch/pyexecutor/sampler/beam_search.py | 29 +++++++++------- .../_torch/sampler/test_beam_search.py | 34 ++++++++++++++++--- 2 files changed, 47 insertions(+), 16 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py index e469dda0e720..f166a8488605 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py @@ -934,20 +934,25 @@ def beam_search_sampling_batch_cba( cba.cba_tokens[slots, :, :snap_len] = merged_tokens cba.cba_log_probs[slots, :, :snap_len] = merged_lps cba.batch_dones[slots] = done - # Publish the done verdict by flooding the full row: the stop criterion - # reads the first py_beam_width (== capacity) entries, which can exceed - # this step's beam_width_out for variable-beam-width requests. + # Publish the done verdict across the full row: the stop criterion reads the + # first py_beam_width (== capacity) entries, which can exceed this step's + # beam_width_out for variable-beam-width requests. # - # Only write where the verdict is done. This tensor is the finish handler's - # first_finish_reasons, which records the *earliest* reason each beam - # finished by and must never be cleared: a beam that ended on a stop word - # vacates its slot and keeps generating on this path, so an unconditional - # write would erase that STOP_WORDS entry on the next step and the request - # would report the reason it eventually stopped for (LENGTH) instead. + # This tensor is the finish handler's first_finish_reasons, which records + # the reason each beam *first* finished by, so only fill entries still at + # NOT_FINISHED. A beam that ended on a stop word does not freeze here -- it + # vacates its slot and the request keeps generating -- so its STOP_WORDS + # entry is the only record of why it ended; overwriting it would make the + # request report whatever stopped it later instead. + # + # Beams with no reason of their own are ending because the pool can no + # longer be beaten, not because they hit a token: report that as LENGTH, + # matching what the pool-free path produced for the same situation. + prev_reasons = args.finished_beams[slots] args.finished_beams[slots] = torch.where( - done.view(-1, 1), - torch.full_like(args.finished_beams[slots], FinishReason.END_ID.value), - args.finished_beams[slots], + done.view(-1, 1) & (prev_reasons == FinishReason.NOT_FINISHED.value), + torch.full_like(prev_reasons, FinishReason.LENGTH.value), + prev_reasons, ) stop_window = args.stop_past_tokens if reordered_window is not None and stop_window is not None: diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 17976369d77b..471a17685288 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -166,7 +166,13 @@ def check_generation_logits(beam: CompletionOutput, """Check if the generation logits have the correct shape""" if sampling_params.return_generation_logits: gen_logits = beam.generation_logits - generated_tokens = valid_tokens if valid_tokens is not None else sampling_params.max_tokens + # Fall back to this beam's own length rather than max_tokens: under + # early_stopping=1 (the default, HF's `True`) the request stops once + # best_of candidates are complete, so a beam that never hit a stop + # token can be shorter than max_tokens. + assert beam.token_ids is not None + generated_tokens = (valid_tokens if valid_tokens is not None else len( + beam.token_ids)) assert gen_logits is not None, "generation logits should not be None" assert gen_logits.ndim == 2, f"generation logits should have 2 dimensions, but got {gen_logits.ndim}" assert gen_logits.shape[ @@ -180,7 +186,13 @@ def check_logprobs(beam: CompletionOutput, sampling_params: SamplingParams, """Check if the logprobs have the correct shape""" assert beam.logprobs is not None if sampling_params.logprobs is not None: - generated_tokens = valid_tokens if valid_tokens is not None else sampling_params.max_tokens + # Fall back to this beam's own length rather than max_tokens: under + # early_stopping=1 (the default, HF's `True`) the request stops once + # best_of candidates are complete, so a beam that never hit a stop + # token can be shorter than max_tokens. + assert beam.token_ids is not None + generated_tokens = (valid_tokens if valid_tokens is not None else len( + beam.token_ids)) assert len( beam.logprobs ) == generated_tokens, f"expected {generated_tokens} logprobs, but got {len(beam.logprobs)}" @@ -206,7 +218,13 @@ def check_cache_indirection(beam: CompletionOutput, assert cache_indirection.shape[ 1] == sampling_params.best_of, f"expected {sampling_params.best_of} entries in dim 1 of cache indirection, but got {cache_indirection.shape[1]}" - num_generated_tokens = valid_tokens if valid_tokens is not None else sampling_params.max_tokens + # Fall back to what this beam actually produced rather than max_tokens: with + # early_stopping=1 (the default, HF's `True`) the request stops as soon as + # best_of finished candidates exist, so a beam that never hit a stop token + # can still be shorter than max_tokens. + assert beam.token_ids is not None + num_generated_tokens = (valid_tokens if valid_tokens is not None else len( + beam.token_ids)) # We return the cache indirection before the sampling step, therefore cache indirection does not reflect changes during the sampling of the last token num_valid_cache_indirection = num_generated_tokens - 1 @@ -250,8 +268,16 @@ def validate_output_beam(beam_output: CompletionOutput, # Check output similarity assert valid_tokens is None or valid_tokens > 0 + # get_expected_outputs walks a fixed number of iterations; it has no notion + # of the finished-candidate pool. Under early_stopping=1 the request stops + # once best_of candidates are complete, so a beam that never hit a stop + # token can end early. Compare the prefix it did produce -- a wrong beam + # still diverges, only the unreached tail is dropped. + assert beam_output.token_ids is not None + num_valid = valid_tokens if valid_tokens is not None else len( + beam_output.token_ids) expected_valid_token_ids = expected_outputs.outputs[ - beam_idx, :valid_tokens].tolist() + beam_idx, :num_valid].tolist() assert beam_output.token_ids == expected_valid_token_ids, f"expected {expected_valid_token_ids} token ids, but got {beam_output.token_ids}" From e86904d54856d5eddf2838e7f66867028ab7c7f7 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Thu, 6 Aug 2026 05:48:30 -0700 Subject: [PATCH 25/36] [TRTLLM-13234][chore] Address beam-search review nits - 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 --- .../_torch/pyexecutor/sampler/beam_search.py | 29 ++++++++++++------- .../_torch/pyexecutor/sampler/sampler.py | 8 ++--- .../host_profile_tools/host_profiler.py | 2 +- .../_torch/sampler/test_beam_search.py | 4 +-- .../_torch/sampler/test_logits_logprobs.py | 2 +- 5 files changed, 25 insertions(+), 20 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py index f166a8488605..fea00c76aa38 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py @@ -77,7 +77,16 @@ def from_raw(cls, value: Optional[int]) -> "BeamSearchEarlyStop": ``None`` -> ``TRUE`` (the default); ``0`` -> ``FALSE``; every other integer -> ``NEVER`` (HF's "never"), matching the CBA path which - special-cases only ``FALSE``.""" + special-cases only ``FALSE``. + + This is deliberately more permissive than the OpenAI-compatible + server, whose ``bool | Literal["never"]`` field rejects anything else + outright. ``sampling_config.early_stopping`` is a plain int that + predates the tri-state and reaches here from the C++ runtime as well, + so values outside {0, 1, 2} are folded into the nearest defined mode + rather than raising on a path that used to accept them. Tightening it + would be an API break for existing callers; the HTTP layer is free to + be strict because it is a newer surface.""" if value is None: return cls.TRUE if value == cls.FALSE: @@ -976,13 +985,14 @@ def beam_search_sampling_batch_cba( args.new_log_probs[slots, :num_beams] = step_log_probs # Record this step's per-slot log-prob at the emission position so path # snapshots can recover per-token log-probs (analog of original_tokens). - olp = cba.original_log_probs[slots, :num_beams] - olp.scatter_( + # Advanced indexing copies, so scatter into the copy and write it back. + beam_log_probs = cba.original_log_probs[slots, :num_beams] + beam_log_probs.scatter_( 2, args.seq_lens.view(-1, 1, 1).expand(-1, num_beams, 1).long(), step_log_probs.unsqueeze(-1), ) - cba.original_log_probs[slots, :num_beams] = olp + cba.original_log_probs[slots, :num_beams] = beam_log_probs args.cum_log_probs[slots, :num_beams] = slot_cum return _pad_next_tokens(slot_tok, args.finished_beams.size(1)), softmax @@ -1018,7 +1028,7 @@ class CBAGroupHost: """None when no request in the group requests log-probs.""" -def _prepare_beam_search( +def prepare_beam_search( beam_search_store: BeamSearchStore, log_probs_store: LogProbsStore, seq_slots_long: torch.Tensor, @@ -1167,7 +1177,7 @@ def _builder() -> BeamHistory | None: return _builder -def _finalize_beam( +def finalize_beam( request: LlmRequest, beam_history: BeamHistory, ) -> None: @@ -1238,8 +1248,6 @@ def __init__( self, *, store: Optional[BeamSearchStore], - log_probs_store: LogProbsStore, - new_tokens: torch.Tensor, max_seq_len: int, max_num_sequences: int, use_speculative_d2h: bool, @@ -1248,8 +1256,6 @@ def __init__( make_side_stream_copier: Callable[[], AbstractContextManager["_SideStreamCopier"]], ): self._store = store - self._log_probs_store = log_probs_store - self._new_tokens = new_tokens self._max_seq_len = max_seq_len self._use_speculative_d2h = use_speculative_d2h # Bound methods of the owning sampler: the D2H copies must share its @@ -1321,7 +1327,8 @@ def prepare_cba_group_host( store = self._store assert store is not None # cba_requests is non-empty here, so the tensors were allocated at - # admission (see BeamSearchHandler.ensure_cba_for_requests). + # admission (see BeamSearchStore.ensure_cba, called from + # TorchSampler._sample_batched_by_strategy). cba = store.cba assert cba is not None, "CBA tensors must be allocated before a CBA step" slots = [request.py_seq_slot for request in cba_requests] diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 2cb0b63131aa..8bc45d5d97b8 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -87,7 +87,7 @@ from ..llm_request import LlmRequest, LlmRequestState, get_draft_token_length from ..resource_manager import ResourceManager, ResourceManagerType from ..scheduler import ScheduledRequests -from .beam_search import BeamHistoryBuilder, BeamSearchHandler, _finalize_beam, _prepare_beam_search +from .beam_search import BeamHistoryBuilder, BeamSearchHandler, finalize_beam, prepare_beam_search from .finish_reasons import FinishReasonsHandler from .logprobs import ( LogProbsState, @@ -1570,8 +1570,6 @@ def __init__(self, args: Args): # speculative predictor reads, so no separate host mirror is kept here. self._beam_search = BeamSearchHandler( store=self.store.beam_search_store, - log_probs_store=self.store.log_probs_store, - new_tokens=self.store.new_tokens, max_seq_len=self.max_seq_len, max_num_sequences=self.max_num_sequences, use_speculative_d2h=self._use_speculative_beam_history_d2h, @@ -2070,7 +2068,7 @@ def setup_sampler_step(self, scheduled_requests: ScheduledRequests) -> None: # Allocate the CBA tensors on the first beam-search request: every # early_stopping mode runs on that path. beam_search_store.ensure_cba() - _prepare_beam_search( + prepare_beam_search( beam_search_store, self.store.log_probs_store, seq_slots_long=seq_slots_tensor_cuda_long, @@ -2426,7 +2424,7 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None: if req.py_beam_width > 1: if (beam_history := _maybe_build_beam_history(req_idx)) is not None: - _finalize_beam(req, beam_history) + finalize_beam(req, beam_history) else: # Only the leading beam_width_out columns hold real tokens; # the op pads the rest of the store-width row with diff --git a/tensorrt_llm/tools/profiler/host_profile_tools/host_profiler.py b/tensorrt_llm/tools/profiler/host_profile_tools/host_profiler.py index 00323ab10d11..02bf6219a180 100644 --- a/tensorrt_llm/tools/profiler/host_profile_tools/host_profiler.py +++ b/tensorrt_llm/tools/profiler/host_profile_tools/host_profiler.py @@ -159,7 +159,7 @@ def resolve(self) -> Optional[Callable]: "update_requests", "_process_requests", "_write_finish_reasons", - "_prepare_beam_search", + "prepare_beam_search", "_select_generated_logits", "_sample_batched_by_strategy", ], diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 471a17685288..f094cef52430 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -40,7 +40,7 @@ SampleStateTorch, TorchSampler) from tensorrt_llm._torch.pyexecutor.sampler.beam_search import ( - CBAGroupHost, _finalize_beam, _gather_beam_path, _prepare_beam_history_cba) + CBAGroupHost, _gather_beam_path, _prepare_beam_history_cba, finalize_beam) from tensorrt_llm._torch.pyexecutor.sampler.sampler_strategy import ( BEAM_SEARCH_PAD_TOKEN, BeamSearch, BeamSearchEarlyStop, BeamSearchMetadata, CBAState, _StrategyImpls, beam_search_sampling_batch_cba) @@ -1942,7 +1942,7 @@ def _uut(): cum_logprobs=cum_logprobs[batch_idx, :beam_width]) request.py_return_log_probs = False - _finalize_beam(request, beam_history) + finalize_beam(request, beam_history) token_history.append(deepcopy(request.get_tokens())) diff --git a/tests/unittest/_torch/sampler/test_logits_logprobs.py b/tests/unittest/_torch/sampler/test_logits_logprobs.py index dc9bc4808d25..d65049fbfc28 100644 --- a/tests/unittest/_torch/sampler/test_logits_logprobs.py +++ b/tests/unittest/_torch/sampler/test_logits_logprobs.py @@ -1362,7 +1362,7 @@ def validate_logprob_and_rank(token_id: int, returned_logprob: Logprob): # There are two factors. First, "rank" is not clearly specified for beam search # (could be logprob rank within beam or across all beams) and therefore # 'rank=1' is returned for all finished beams - # via _finalize_beam and convert_logprobs_tensor_to_list. Second, + # via finalize_beam and convert_logprobs_tensor_to_list. Second, # during decoding (unfinished beam, in general this is the case in this test), # request logprobs are set via handle_logprobs and store_logprobs_list_to_request, # which inspects uninitialized elements of log_probs_store.sampled_log_prob_ranks. From 3ec722544cc9e041bc6a0e65e4e415cb1ea670d8 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Thu, 6 Aug 2026 06:10:12 -0700 Subject: [PATCH 26/36] [TRTLLM-13234][fix] Give the top-p-decay mock request a beam width 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 --- tests/unittest/_torch/sampler/test_torch_sampler.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index 4a29f602bba6..d38ac2f18076 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -3455,6 +3455,9 @@ def _mock_request(params: SamplingParams, *, draft_tokens=None): is_context_init_state=False, py_sampling_strategy=None, py_draft_tokens=draft_tokens, + # Read by the row_stride query in sampler_common; these tests are + # single-beam, so the static admission width is 1. + py_beam_width=1, ) req.get_beam_width_by_iter = lambda for_next_iteration=False: 1 return cast(LlmRequest, req) From 6e975683de6c105f74aff07dc388c863c453cc38 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Thu, 6 Aug 2026 06:53:34 -0700 Subject: [PATCH 27/36] [TRTLLM-13234][fix] Expect LENGTH from the CBA done verdict in its unit 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 --- tests/unittest/_torch/sampler/test_beam_search.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index f094cef52430..a75bcfc0bb6a 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -1243,7 +1243,9 @@ def test_beam_search_cba_done_bound_by_early_stopping(early_stopping, return_probs=False, ) assert m.cba.batch_dones[0].item() is expect_done - expected_reason = (FinishReason.END_ID.value + # A beam with no reason of its own ends because the pool can no longer be + # beaten, not because it hit a token, so the verdict publishes LENGTH. + expected_reason = (FinishReason.LENGTH.value if expect_done else FinishReason.NOT_FINISHED.value) assert (m.finished_beams[0] == expected_reason).all() # CBA untouched either way (no eligible end candidates this step). From d84c778f3dc55d735ca7d53848358402db00abd2 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Thu, 6 Aug 2026 19:14:22 -0700 Subject: [PATCH 28/36] [TRTLLM-13234][fix] Give the remaining request mocks a beam width 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 --- tests/unittest/_torch/sampler/test_torch_sampler.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index d38ac2f18076..b9dcd85d195a 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -146,6 +146,9 @@ class MockLlmRequest: sampling_config: SamplingConfig is_context_init_state: bool # Torch sampler accesses this, but it does not affect this test py_sampling_strategy: Strategy | None + # Read by the row_stride query in sampler_common; these tests are + # single-beam, so the static admission width is 1. + py_beam_width: int def get_beam_width_by_iter( self, for_next_iteration: bool = False @@ -169,6 +172,7 @@ def _build_mock_llm_request(self, params: SamplingParams) -> LlmRequest: request.sampling_config = SamplingConfig(params._get_sampling_config()) request.is_context_init_state = False # Not used in this test request.py_sampling_strategy = None # used for caching + request.py_beam_width = 1 return cast(LlmRequest, request) def test_defaults(self): @@ -561,6 +565,8 @@ class GenRequestMock: def __init__(self, draft_len: int): self.py_draft_tokens = torch.empty(draft_len, dtype=torch.int32, device=device) self.sampling_config = SamplingConfig(beam_width=1) + # Read by the row_stride query in sampler_common. + self.py_beam_width = 1 def get_beam_width_by_iter( self, for_next_iteration: bool = False From e2bf7055131494e010ec080e4df05d8d8173eec6 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Thu, 6 Aug 2026 19:30:56 -0700 Subject: [PATCH 29/36] [TRTLLM-13234][fix] Give the logprobs tests Dynamo recompile headroom 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 --- .../_torch/sampler/test_logits_logprobs.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/unittest/_torch/sampler/test_logits_logprobs.py b/tests/unittest/_torch/sampler/test_logits_logprobs.py index d65049fbfc28..9637d8b47c8c 100644 --- a/tests/unittest/_torch/sampler/test_logits_logprobs.py +++ b/tests/unittest/_torch/sampler/test_logits_logprobs.py @@ -38,6 +38,22 @@ ) +@pytest.fixture(autouse=True) +def _dynamo_recompile_headroom(): + """Recompile headroom for the fullgraph=True beam-search step. + + The beam-search cases here build a fresh engine per parametrization in one + process, and Dynamo counts recompiles per code object across the process, + so the default limit (8) is exhausted partway through and the rest hard-fail + under fullgraph. The limit guards against runaway recompilation and is not a + correctness property; a served model has a fixed set of shapes. + """ + import torch._dynamo + + with torch._dynamo.config.patch(recompile_limit=128): + yield + + @pytest.fixture(scope="module", params=[False, True]) def disable_overlap_scheduler_fixture(request) -> bool: return request.param From b11b92870757da497115375a32fa0046acb3d4ce Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Thu, 6 Aug 2026 19:47:33 -0700 Subject: [PATCH 30/36] [TRTLLM-13234][fix] Give the compiled CBA step its own recompile headroom 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 --- .../_torch/pyexecutor/sampler/beam_search.py | 20 +++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py index fea00c76aa38..abe65082e05a 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py @@ -23,7 +23,7 @@ from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass from enum import IntEnum -from typing import TYPE_CHECKING, Callable, NamedTuple, Optional, TypeAlias, cast +from typing import TYPE_CHECKING, Any, Callable, NamedTuple, Optional, TypeAlias, cast import torch @@ -751,7 +751,23 @@ def _take(src: torch.Tensor) -> torch.Tensor: # Compiled lazily on first use; the first CBA-mode request of a process pays # the inductor compile once (subsequent shapes are covered by the dynamic # batch/snapshot-length dims marked at the call site). -_cba_step_compiled = torch.compile(_cba_step_math, dynamic=None, fullgraph=True) +# +# Dynamo counts recompiles per code object for the lifetime of the process, and +# fullgraph turns exhaustion into a hard failure rather than a fallback to +# eager. A served engine sees one shape family and never approaches the default +# cap of 8, but a process that builds many engines -- a test session, or a +# worker reused across configurations -- does. Give this one function its own +# headroom so a later engine in the same process still compiles; the cap exists +# to catch runaway recompilation, not as a correctness property. +_CBA_RECOMPILE_LIMIT = 256 + + +def _cba_step_compiled(*args: Any, **kwargs: Any) -> CBAStepResult: + with torch._dynamo.config.patch(recompile_limit=_CBA_RECOMPILE_LIMIT): + return _cba_step_compiled_inner(*args, **kwargs) + + +_cba_step_compiled_inner = torch.compile(_cba_step_math, dynamic=None, fullgraph=True) def beam_search_sampling_batch_cba( From acf489cdd60c976d0a6a74b11b15e5b282f5efb6 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Thu, 6 Aug 2026 21:07:31 -0700 Subject: [PATCH 31/36] [TRTLLM-13234][fix] Latch the CBA harvest separately from the finish 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 --- .../_torch/pyexecutor/sampler/beam_search.py | 25 +++++- .../pyexecutor/sampler/finish_reasons.py | 12 +++ .../_torch/pyexecutor/sampler/sampler.py | 4 + .../_torch/sampler/test_beam_search.py | 90 +++++++++++++++++++ 4 files changed, 127 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py index abe65082e05a..65e733b62d99 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py @@ -208,6 +208,16 @@ class BeamSearchStore: of each active beam.""" first_finish_reasons: torch.Tensor """[max_num_sequences, max_beam_width] int32, first finish reason per beam.""" + pending_harvest: torch.Tensor + """[max_num_sequences, max_beam_width] bool, beams latched finished by the + finish handler and not yet harvested into the CBA. + + Kept apart from ``first_finish_reasons`` on purpose. That tensor is the + request's reported finish reason and must survive for the whole request, + while this one is a one-shot signal: the harvest consumes it, and the slot + it frees goes on to hold an unrelated, unfinished continuation. Reading the + persistent reason as the latch would harvest that continuation again on the + next step.""" predecessor_beams: torch.Tensor """[max_num_sequences, max_beam_width] int32, predecessor beam per beam, used for stop word detection.""" @@ -248,6 +258,7 @@ def create( predecessor_beams=int_tensor(per_beam), original_tokens=int_tensor(cache_indirection_shape), first_finish_reasons=int_tensor(per_beam), + pending_harvest=torch.zeros(per_beam, device="cuda", dtype=torch.bool), beam_idx_arange=torch.arange(max_beam_width, device="cuda", dtype=torch.int32), prompt_lens=int_tensor((max_num_sequences,)), batch_dones=torch.zeros((max_num_sequences,), device="cuda", dtype=torch.bool), @@ -312,6 +323,7 @@ class BeamSearchMetadata(StrategyMetadata): seq_slots: torch.Tensor seq_lens: torch.Tensor finished_beams: torch.Tensor + pending_harvest: torch.Tensor predecessor_beams: torch.Tensor beam_idx_arange: torch.Tensor stop_past_tokens: Optional[torch.Tensor] = None @@ -548,7 +560,7 @@ def _cba_step_math( original_tokens: torch.Tensor, original_log_probs: torch.Tensor, cum_log_probs: torch.Tensor, - finished_beams: torch.Tensor, + pending_harvest: torch.Tensor, end_ids: torch.Tensor, prompt_lens: torch.Tensor, cba_caps: torch.Tensor, @@ -591,7 +603,7 @@ def _cba_step_math( batch_size, num_candidates = cand_cum.shape neg_inf = float("-inf") - harvest_mask = finished_beams[slots, :beam_width_in] != FinishReason.NOT_FINISHED.value + harvest_mask = pending_harvest[slots, :beam_width_in] end_ids_b = end_ids[slots].view(-1, 1) caps = cba_caps[slots].view(-1, 1) prompts = prompt_lens[slots].view(-1, 1) @@ -851,7 +863,7 @@ def beam_search_sampling_batch_cba( # lengths are uniform here) and the END_ID flood below is all-beams too — # either way the request stops that same step, so there is no next step # to harvest them (finalize handles those paths instead). - harvest_mask = args.finished_beams[slots, :beam_width_in] != FinishReason.NOT_FINISHED.value + harvest_mask = args.pending_harvest[slots, :beam_width_in] logprobs = logprobs.masked_fill(harvest_mask.unsqueeze(-1), float("-inf")) logprobs += args.cum_log_probs.unsqueeze(-1)[slots, :beam_width_in] @@ -936,7 +948,7 @@ def beam_search_sampling_batch_cba( original_tokens=cba.original_tokens, original_log_probs=cba.original_log_probs, cum_log_probs=args.cum_log_probs, - finished_beams=args.finished_beams, + pending_harvest=args.pending_harvest, end_ids=cba.end_ids, prompt_lens=cba.prompt_lens, cba_caps=cba.cba_caps, @@ -959,6 +971,10 @@ def beam_search_sampling_batch_cba( cba.cba_tokens[slots, :, :snap_len] = merged_tokens cba.cba_log_probs[slots, :, :snap_len] = merged_lps cba.batch_dones[slots] = done + # The latch is one-shot: the beams it named have just been pooled and their + # slots refilled with unrelated continuations, so leaving it set would + # harvest those on the next step (see BeamSearchStore.pending_harvest). + args.pending_harvest[slots] = torch.zeros_like(args.pending_harvest[slots]) # Publish the done verdict across the full row: the stop criterion reads the # first py_beam_width (== capacity) entries, which can exceed this step's # beam_width_out for variable-beam-width requests. @@ -1069,6 +1085,7 @@ def prepare_beam_search( beam_search_store.first_finish_reasons.index_fill_( 0, seq_slots_long, FinishReason.NOT_FINISHED.value ) + beam_search_store.pending_harvest.index_fill_(0, seq_slots_long, False) beam_search_store.original_tokens.index_fill_(0, seq_slots_long, 0) beam_search_store.prompt_lens.index_copy_(0, seq_slots_long, prompt_lens_cuda) beam_search_store.batch_dones.index_fill_(0, seq_slots_long, False) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py b/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py index e79281176391..a4a818541b4b 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py @@ -528,6 +528,7 @@ def write_finish_reasons( seq_lens_cuda: torch.Tensor, new_tokens_cuda: torch.Tensor, first_finish_reasons_cuda: torch.Tensor | None = None, + pending_harvest_cuda: torch.Tensor | None = None, ) -> torch.Tensor: """Calculates the finish reasons for each request and returns the finish reasons tensor. @@ -567,6 +568,7 @@ def write_finish_reasons( stop_word_indices=stop_word_indices_cuda, single_token_stop_words_only=single_token_stop_words_only, first_finish_reasons=first_finish_reasons_cuda, + pending_harvest=pending_harvest_cuda, ) return self.store.finish_reasons_cuda @@ -638,6 +640,7 @@ def _write_finish_reasons( stop_word_indices: torch.Tensor | None = None, single_token_stop_words_only: bool = False, first_finish_reasons: torch.Tensor | None = None, + pending_harvest: torch.Tensor | None = None, ) -> None: """Writes the finish reasons to the finish_reasons tensor. @@ -712,11 +715,20 @@ def _write_finish_reasons( if first_finish_reasons is not None: # store the first stop reason for each beam of a seq_slot. batched_first_finish_reasons = first_finish_reasons[seq_slots] + newly_finished = (batched_first_finish_reasons == FinishReason.NOT_FINISHED.value) & ( + batched_finish_reasons != FinishReason.NOT_FINISHED.value + ) first_finish_reasons[seq_slots, ...] = torch.where( batched_first_finish_reasons == FinishReason.NOT_FINISHED.value, batched_finish_reasons, batched_first_finish_reasons, ) + if pending_harvest is not None: + # Raise the beam-search harvest latch for beams that finished on + # *this* step. The CBA step lowers it once it has pooled them; + # first_finish_reasons itself cannot serve as the latch because + # it must outlive the harvest to be reported to the caller. + pending_harvest[seq_slots, ...] |= newly_finished.any(dim=0) def _are_end_id(self, end_ids_cuda: torch.Tensor, tokens_cuda: torch.Tensor) -> torch.Tensor: """Checks if the tokens are the end id diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 8bc45d5d97b8..5469c3552fb8 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -2255,6 +2255,7 @@ def _add_metadata_to_grouped_requests( seq_slots=group_seq_slots_cuda, seq_lens=group_seq_lens_cuda, finished_beams=beam_search_store.first_finish_reasons, + pending_harvest=beam_search_store.pending_harvest, predecessor_beams=beam_search_store.predecessor_beams, beam_idx_arange=beam_search_store.beam_idx_arange, stop_past_tokens=self._finish_reasons_handler.store.past_tokens_cuda, @@ -2658,6 +2659,9 @@ def sample_async( if beam_search_store is not None else None ), + pending_harvest_cuda=( + beam_search_store.pending_harvest if beam_search_store is not None else None + ), ) finish_reasons_host = self._copy_to_host(finish_reasons_device) diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index a75bcfc0bb6a..5c31eed18227 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -989,6 +989,8 @@ def make_metadata() -> BeamSearchMetadata: seq_lens=torch.full((batch_size, ), seq_len, dtype=torch.int32), finished_beams=torch.zeros((max_batch_size, beam_width), dtype=torch.int32), + pending_harvest=torch.zeros((max_batch_size, beam_width), + dtype=torch.bool), new_log_probs=torch.zeros((max_batch_size, beam_width), dtype=torch.float32), predecessor_beams=torch.zeros((max_batch_size, beam_width), @@ -1096,6 +1098,9 @@ def _make_cba_metadata(max_batch, K, attn_len, snap_len, seq_len, prompt_len, seq_slots=slots, seq_lens=torch.full((batch, ), seq_len, dtype=torch.int32), finished_beams=torch.zeros((max_batch, K), dtype=torch.int32), + # One-shot latch the finish handler raises and the CBA step consumes; + # tests set it alongside finished_beams to stage a pending harvest. + pending_harvest=torch.zeros((max_batch, K), dtype=torch.bool), predecessor_beams=torch.zeros((max_batch, K), dtype=torch.int32), beam_idx_arange=torch.arange(K, dtype=torch.int32), cba=CBAState( @@ -1418,6 +1423,7 @@ def test_beam_search_cba_harvest_stop_word_beam(): m.cum_log_probs[0] = torch.tensor([-1.5, -2.0]) # beam 0 was latched STOP_WORDS by the finish handler after last step m.finished_beams[0, 0] = FinishReason.STOP_WORDS.value + m.pending_harvest[0, 0] = True logits = torch.full((K, vocab), -50.0) logits[0, 1] = 10.0 # beam0's candidates must be ignored (harvested) @@ -1450,6 +1456,90 @@ def test_beam_search_cba_harvest_stop_word_beam(): assert tokens[0].tolist() == [2, 3] +@_kernel_test +def test_beam_search_cba_harvest_latch_clears_after_refill(): + """A harvested slot must not be harvested again on the next step. + + The harvest mask reads first_finish_reasons, which is also the request's + persistent output finish reason and so survives the refill. Reading it as + a transient latch means the *new* continuation occupying that slot looks + finished on the following step: it gets harvested a second time, the pool + gains an entry for a hypothesis that never ended, and the request can be + declared done early. + """ + from tensorrt_llm._torch.pyexecutor.sampler.beam_search import \ + beam_search_sampling_batch_cba + + K, vocab, end_id = 2, 6, 5 + prompt, gen = 2, 3 + seq_len = prompt + gen + m = _make_cba_metadata(max_batch=1, + K=K, + attn_len=10, + snap_len=6, + seq_len=seq_len, + prompt_len=prompt, + end_id=end_id, + batch=1) + for b in range(K): + m.cache_indirection[0, b, :] = b + m.cba.original_tokens[0, b, prompt:seq_len] = torch.tensor( + [70 + b, 71 + b, 72 + b], dtype=torch.int32) + m.cum_log_probs[0] = torch.tensor([-1.5, -2.0]) + # Step 1: beam 0 was latched STOP_WORDS by the finish handler. + m.finished_beams[0, 0] = FinishReason.STOP_WORDS.value + m.pending_harvest[0, 0] = True + + logits = torch.full((K, vocab), -50.0) + logits[0, 1] = 10.0 # ignored: beam 0 is harvested, not expanded + logits[1, 2] = 9.0 + logits[1, 3] = 8.0 + + beam_search_sampling_batch_cba( + logits=logits, + beam_width_in=K, + beam_width_out=K, + beam_search_args=m, + temperature=1.0, + early_stopping=0, + length_penalty=1.0, + return_probs=False, + ) + # Step 2 verifies the harvest happened and both slots refilled from beam 1. + entries_after_first = int((m.cba.cba_normed_scores[0] + > _CBA_NEG_INF).sum().item()) + assert entries_after_first == 1, "the stop-word beam should be pooled once" + assert m.predecessor_beams[0].tolist() == [1, 1] + + # Step 3: run again. Neither slot finished this step -- they hold the + # continuations that refilled them, and the finish handler latched nothing + # new, so nothing may be harvested. + m.seq_lens = torch.full((1, ), seq_len + 1, dtype=torch.int32) + logits2 = torch.full((K, vocab), -50.0) + logits2[0, 1] = 7.0 + logits2[1, 2] = 6.0 + + beam_search_sampling_batch_cba( + logits=logits2, + beam_width_in=K, + beam_width_out=K, + beam_search_args=m, + temperature=1.0, + early_stopping=0, + length_penalty=1.0, + return_probs=False, + ) + # Step 4: the refilled continuation must not be pooled or masked again. + entries_after_second = int((m.cba.cba_normed_scores[0] + > _CBA_NEG_INF).sum().item()) + assert entries_after_second == entries_after_first, ( + f"pool grew from {entries_after_first} to {entries_after_second}: an " + "unfinished continuation was harvested because the stop-word latch " + "outlived the refill") + assert not m.cba.batch_dones[0].item(), ( + "request finished early: the re-harvest emptied the beam slots") + + @_kernel_test def test_beam_search_cba_reorders_stop_window(): """The finish handler's stop-word window must follow beam swaps.""" From 44612d6dfc26f834a9d4368135950f81affb862c Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Fri, 7 Aug 2026 02:06:52 -0700 Subject: [PATCH 32/36] [TRTLLM-13234][fix] Annotate two returns mypy could not infer 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 --- tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py | 3 ++- tests/unittest/_torch/sampler/test_beam_search.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py index af265b33f987..ae43820dd1df 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py @@ -82,7 +82,8 @@ def radix_topk_op( order among equal values vary run to run. This op sits on the default beam path, where reproducible tie order matters for triage. """ - return flashinfer.top_k(values, k, sorted=True, deterministic=True) + topk_values, topk_indices = flashinfer.top_k(values, k, sorted=True, deterministic=True) + return topk_values, topk_indices @_compiler_disable diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 5c31eed18227..2481214ce19f 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -795,7 +795,7 @@ def _kernel_test(fn: Callable[..., Any]) -> Callable[..., Any]: @pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") @functools.wraps(fn) - def wrapper(*args, **kwargs): + def wrapper(*args: Any, **kwargs: Any) -> Any: with torch.device("cuda"): return fn(*args, **kwargs) From 8da271eba65ec70cbbbb8f5c6547b33ef1b93b7c Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Fri, 7 Aug 2026 20:35:45 -0700 Subject: [PATCH 33/36] [TRTLLM-13234][fix] Name _kernel_test's return type past functools.wraps 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 --- tests/unittest/_torch/sampler/test_beam_search.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 2481214ce19f..6238deb01903 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -799,7 +799,10 @@ def wrapper(*args: Any, **kwargs: Any) -> Any: with torch.device("cuda"): return fn(*args, **kwargs) - return wrapper + # functools.wraps is untyped, so `wrapper` decays to Any; name the type + # again on the way out to satisfy the declared return. + typed_wrapper: Callable[..., Any] = wrapper + return typed_wrapper class GeneralTestParams: From b60085581816cca9f144a2d4514ac56dc51513d5 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Sat, 8 Aug 2026 03:46:37 -0700 Subject: [PATCH 34/36] [TRTLLM-13234][fix] Detect syncs without parking a hostfunc on the stream 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 --- .../test_beam_search_speculative_d2h.py | 44 ++++++++++++++++--- 1 file changed, 37 insertions(+), 7 deletions(-) diff --git a/tests/unittest/_torch/sampler/test_beam_search_speculative_d2h.py b/tests/unittest/_torch/sampler/test_beam_search_speculative_d2h.py index 3a5004f30aee..fd3a38b53257 100644 --- a/tests/unittest/_torch/sampler/test_beam_search_speculative_d2h.py +++ b/tests/unittest/_torch/sampler/test_beam_search_speculative_d2h.py @@ -30,15 +30,16 @@ """ import gc +import os import pathlib as _pl -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext from copy import deepcopy -from typing import Any, Iterable +from typing import Any, Generator, Iterable import pytest +import torch from pydantic import ValidationError from test_beam_search_util import DummyConfigLoader, DummyWeightLoader -from utils.util import assert_no_cuda_sync from tensorrt_llm import LLM, SamplingParams from tensorrt_llm._torch.models.checkpoints import HfCheckpointLoader @@ -421,6 +422,35 @@ def test_speculative_d2h_rejects_async_worker_combo( # --------------------------------------------------------------------------- +@contextmanager +def _assert_no_cuda_sync_locally() -> Generator[None, None, None]: + """Local stand-in for utils.util.assert_no_cuda_sync. + + That helper parks a @hostfunc on the stream to block it, and only lowers + its cancel flag *after* the assertion. The hostfunc holds the GIL while it + spins, so anything that blocks before the cancel deadlocks the step: the + main thread waits on the stream, the stream waits on the hostfunc, and the + hostfunc waits for a cancel that the main thread never reaches. On the + speculative path -- side-stream copier plus the executor's own threads -- + that race is live, and CI hangs here until the 3600s pytest timeout. + + Detect the same thing without parking anything on the stream: + set_sync_debug_mode("error") makes torch raise on a synchronizing call, + which is the actual property under test. It misses syncs from non-torch + kernels, which the stream-blocking variant would catch; nothing in + sample_async/update_requests issues those today. + """ + if int(os.environ.get("CUDA_LAUNCH_BLOCKING", 0)): + yield None + return + previous_mode = torch.cuda.get_sync_debug_mode() + torch.cuda.set_sync_debug_mode("error") + try: + yield None + finally: + torch.cuda.set_sync_debug_mode(previous_mode) + + @pytest.mark.threadleak(enabled=False) def test_speculative_d2h_predictor_hit_is_sync_free( fixed_params: dict[str, Any], @@ -431,7 +461,7 @@ def test_speculative_d2h_predictor_hit_is_sync_free( inside `sample_async` or inside `update_requests` after the sampler event has been awaited. - Mirrors the `assert_no_cuda_sync()` hook used by + Mirrors the sync check used by `tests/unittest/_torch/sampler/test_beam_search.py::validate_outputs`, but pins the predictor to always-hit so every step routes through the side-stream copier and never reaches the `.cpu()` fallback. @@ -444,17 +474,17 @@ def test_speculative_d2h_predictor_hit_is_sync_free( def _sample_async_hook(self, *args, **kwargs): # type: ignore[no-untyped-def] hook_state["sample_async_called"] = True - with assert_no_cuda_sync(): + with _assert_no_cuda_sync_locally(): return sample_async_orig(self, *args, **kwargs) def _update_requests_hook(self, state: SampleStateTorch, *args, **kwargs): # type: ignore[no-untyped-def] hook_state["update_requests_called"] = True # Sampler event awaits all device work (incl. side-stream copies) - # and is the one expected sync; do it outside assert_no_cuda_sync. + # and is the one expected sync; do it outside the guard below. sampler_event = state.sampler_event if sampler_event: sampler_event.synchronize() - with assert_no_cuda_sync(): + with _assert_no_cuda_sync_locally(): state.sampler_event = None try: return update_requests_orig(self, state, *args, **kwargs) From 238c420ff0e5195333c690178f35a19864025ca7 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Mon, 10 Aug 2026 07:33:24 -0700 Subject: [PATCH 35/36] [TRTLLM-13234][fix] Complete the disaggregated beam search handoff 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 --- tensorrt_llm/_torch/pyexecutor/py_executor.py | 61 +++++++++---------- .../_torch/pyexecutor/sampler/sampler.py | 41 ++++++++++++- .../_torch/sampler/test_beam_search.py | 51 ++++++++-------- 3 files changed, 93 insertions(+), 60 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index ed06c7449733..9dd6c2bbeb4a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -5026,33 +5026,6 @@ def _validate_request(self, request: LlmRequest): f"length_penalty {invalid} is negative; only " "non-negative values are supported.") - # Beam search keeps a pool of finished candidates, which the - # context server can already populate on its first step. That pool - # is not part of the disaggregated 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 -- reachable under aggregation but not under - # disaggregation. - # - # This applies to every early_stopping mode and to both samplers: - # TorchSampler now serves all modes from the candidate-beams-array - # path, and the C++ decoder behind TRTLLMSampler maintains the pool - # unconditionally (beamSearchLayer.cu assigns numBeamsCBA with no - # mode check). Reject until the handoff carries the pool; - # TRTLLM-14792. - # NB: is_context_only_request is a property, but - # is_generation_only_request is a plain method -- it must be called, - # otherwise the bound method is truthy and this matches every request. - if (request.is_context_only_request - or request.is_generation_only_request()): - if sampling_config.beam_width > 1: - raise ValueError( - "Beam search is not supported with disaggregated " - "serving: the finished-candidate pool built during the " - "context phase is not transferred to the generation " - "server, so completions found there would be silently " - "dropped. Use beam_width=1.") - # Check token ID ranges self._validate_token_id_range(request) @@ -6430,12 +6403,13 @@ def _update_sampler_state_for_disagg_gen_request(self, req, beam_width, first_gen_tokens) -> bool: """Update beam sampler state with context-side first-token data. - NB: currently unreachable past the beam_width <= 1 guard. Admission - rejects beam search under disaggregated serving (see - _validate_request), because the finished-candidate pool the context - server builds is not part of the handoff. Kept rather than deleted: - this is the seeding half of that handoff and becomes live again once - the pool is transferred; TRTLLM-14792. + Seeds this side's beam state from the handoff: the per-beam token, its + cumulative log-prob, and an identity cache indirection. A beam whose + token is the end id finished during prefill and is latched for harvest + so the first CBA step pools it here -- the context server's own pool is + not transferred (TRTLLM-14792), which is why its end id is masked for + the CBA op so such a token stays in the beam slot and survives the + handoff. """ if beam_width <= 1: return True @@ -6516,6 +6490,26 @@ def fail_request(message: str) -> bool: # at index prompt_len. No generated-length counter needs seeding to # match: length_penalty normalizes by seq_lens - prompt_lens, which # counts that token because it occupies the first generated slot. + + # A beam whose handed-off token is the end id finished during prefill. + # The context server keeps such a token in its beam slot rather than + # pooling it (its end id is masked for the CBA op, see + # _group_requests_with_metadata), precisely so it survives the handoff + # and can be pooled here instead. Raise the harvest latch and the first + # CBA step folds the path into this side's pool and frees the slot -- + # the same route stop words take. + end_id = req.py_end_id + if end_id is not None and end_id >= 0: + finished_beams = [ + beam_idx for beam_idx in range(beam_width) + if first_gen_tokens[beam_idx] == end_id + ] + if finished_beams: + pending_harvest = beam_search_store.pending_harvest + pending_harvest[seq_slot, + torch.tensor(finished_beams, + device=pending_harvest.device, + dtype=torch.long)] = True return True @staticmethod @@ -6966,6 +6960,7 @@ def _sample_async(self, scheduled_batch, self._handle_errors(error_msg) @nvtx_range("_setup_sampler_step") + @torch.inference_mode() def _setup_sampler_step(self, requests: ScheduledRequests): try: return self.sampler.setup_sampler_step(requests) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 5469c3552fb8..5546b473a3ad 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -2247,6 +2247,27 @@ def _add_metadata_to_grouped_requests( group_seq_lens_cuda = seq_lens[value.indices].to( device="cuda", non_blocking=True ) # Should be on device for beam search + # Context-only (disaggregated prefill) requests hand off after + # this single step, so the finished-candidate pool they would + # build is discarded rather than transferred. Mask their end id + # to the "no end token" sentinel (< 0) so an end candidate stays + # in its beam slot instead: the slot's token is what reaches the + # generation server as first_gen_tokens, which lets that side + # rebuild the pool entry itself. Data-only -- the op sees a + # tensor, so this adds no branch to the compiled step. + cba_end_ids = self._finish_reasons_handler.store.end_ids_cuda + if beam_search_store.cba is not None and any( + requests[i].is_context_only_request for i in value.indices.tolist() + ): + cba_end_ids = cba_end_ids.clone() + ctx_slots = [ + requests[i].py_seq_slot + for i in value.indices.tolist() + if requests[i].is_context_only_request + ] + cba_end_ids[ + torch.tensor(ctx_slots, device=cba_end_ids.device, dtype=torch.long) + ] = -1 metadata = BeamSearchMetadata( cache_indirection=beam_search_store.cache_indirection, cache_indirection_buffer=beam_search_store.cache_indirection_buffer, @@ -2264,7 +2285,7 @@ def _add_metadata_to_grouped_requests( cba=None if beam_search_store.cba is None else CBAState( - end_ids=self._finish_reasons_handler.store.end_ids_cuda, + end_ids=cba_end_ids, prompt_lens=beam_search_store.prompt_lens, original_tokens=beam_search_store.original_tokens, batch_dones=beam_search_store.batch_dones, @@ -2424,7 +2445,23 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None: continue if req.py_beam_width > 1: - if (beam_history := _maybe_build_beam_history(req_idx)) is not None: + # Context-only (disaggregated prefill) requests reach the + # sampler already flagged as finished -- the context server is + # done with them -- but they are migrating, not completing: + # the generation server decodes the rest. Finalizing here + # would rewrite the beam rows from the beam history, and after + # a single step only beam 0 has history; every other beam is + # all BEAM_SEARCH_PAD_TOKEN, so set_generated_tokens would + # give it an empty generated sequence. The handoff reads the + # per-beam first token back out via getTokens().back() + # (llmRequest.cpp), which would then hand the generation + # server the prompt tail instead of that beam's token. Append + # this step's tokens instead and let the generation side own + # finalization. + if ( + not req.is_context_only_request + and (beam_history := _maybe_build_beam_history(req_idx)) is not None + ): finalize_beam(req, beam_history) else: # Only the leading beam_width_out columns hold real tokens; diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 6238deb01903..5b04b8ddbe78 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -448,17 +448,14 @@ def test_beam_search_disagg_e2e( input_prompts, model_kwargs: dict[str, Any], ) -> None: - """Beam search must be rejected under disaggregated serving. - - Every early_stopping mode keeps a pool of finished candidates, which the - context server can already populate on its first step. That pool 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 holds for both - samplers: TorchSampler serves every mode from the candidate-beams-array - path, and the C++ decoder behind TRTLLMSampler maintains the pool - unconditionally. Rejected at admission until the handoff carries the pool - (TRTLLM-14792). + """Beam search is admitted under disaggregated serving. + + The context server's finished-candidate pool is not part of the handoff + (TRTLLM-14792), so the CBA op runs with that side's end id masked: an end + candidate stays in its beam slot rather than being pooled, travels as + first_gen_tokens, and the generation server pools it there instead. Every + early_stopping mode goes through the same route, so admission accepts them + all rather than rejecting beam search outright. """ sampling_params = SamplingParams( max_tokens=fixed_params["max_tokens"], @@ -489,24 +486,28 @@ def test_beam_search_disagg_e2e( ctx_llm = _build_llm(fixed_params, prompts, disagg_kwargs) try: with ctx_llm: - # Every mode is rejected, including the default (early_stopping - # unset, i.e. True). + # Every mode goes through the same route, including the default + # (early_stopping unset, i.e. True). for early_stopping in (None, 0, 1, 2): params = deepcopy(sampling_params) if early_stopping is not None: params.early_stopping = early_stopping - with pytest.raises(RequestError, - match=".*not supported with disaggregated" - " serving.*"): - _ = ctx_llm.generate( - deepcopy(prompts), - sampling_params=params, - disaggregated_params=[ - DisaggregatedParams(request_type="context_only", - disagg_request_id=200) - ], - use_tqdm=False, - ) + outputs = ctx_llm.generate( + deepcopy(prompts), + sampling_params=params, + disaggregated_params=[ + DisaggregatedParams(request_type="context_only", + disagg_request_id=200) + ], + use_tqdm=False, + ) + # The context phase hands off one token per beam; admission no + # longer rejects it, which is what this pins. + assert len(outputs) == len(prompts) + ctx_params = outputs[0].disaggregated_params + assert ctx_params is not None + assert len(ctx_params.first_gen_tokens + ) == fixed_params["max_beam_width"] finally: ctx_llm.shutdown() From 5f4906a23a370248f694be187cc0dc3268faa7c9 Mon Sep 17 00:00:00 2001 From: ZhaoyangWang Date: Mon, 10 Aug 2026 23:34:07 -0700 Subject: [PATCH 36/36] [TRTLLM-13234][fix] Hand off a context phase that ends on its only token 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 --- .../pyexecutor/sampler/finish_reasons.py | 18 +++- .../_torch/pyexecutor/sampler/sampler.py | 24 +---- .../_torch/sampler/test_beam_search.py | 101 ++++++++++++++++++ 3 files changed, 122 insertions(+), 21 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py b/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py index a4a818541b4b..a756113dc4e7 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py @@ -262,7 +262,23 @@ def prepare_for_new_request(self, request: LlmRequest) -> None: self._temp_data.max_lens.append( min(self._max_seq_len, request.orig_prompt_len + request.py_max_new_tokens) ) - self._temp_data.end_ids.append(end_id if (end_id := request.py_end_id) is not None else -1) + # Beam-search context-only (disaggregated prefill) requests hand off + # after their single step, so their end id is masked to the "no end + # token" sentinel (< 0). Two things depend on it: an end candidate is + # not pooled by the CBA op, so it stays in its beam slot and travels to + # the generation server as first_gen_tokens; and the request is not + # marked finished here, so it still reaches the disagg-transmission + # state that builds ContextPhaseParams (llmRequest.cpp) -- an END_ID + # finish leaves it in GENERATION_COMPLETE, which is neither + # isContextFinished() nor finished-due-to-length, so the handoff is + # never started and no tokens are produced at all. The generation + # server sees the end id itself and pools the beam there + # (TRTLLM-14792). Scoped to beam search: single-beam disaggregation + # keeps its current end-id behaviour. + end_id = request.py_end_id + if end_id is None or (request.is_context_only_request and request.py_beam_width > 1): + end_id = -1 + self._temp_data.end_ids.append(end_id) if (stop_words_list := request.py_stop_words_list) is not None: assert (seq_slot := request.py_seq_slot) is not None diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 5546b473a3ad..2597174fb6d8 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -2247,27 +2247,11 @@ def _add_metadata_to_grouped_requests( group_seq_lens_cuda = seq_lens[value.indices].to( device="cuda", non_blocking=True ) # Should be on device for beam search - # Context-only (disaggregated prefill) requests hand off after - # this single step, so the finished-candidate pool they would - # build is discarded rather than transferred. Mask their end id - # to the "no end token" sentinel (< 0) so an end candidate stays - # in its beam slot instead: the slot's token is what reaches the - # generation server as first_gen_tokens, which lets that side - # rebuild the pool entry itself. Data-only -- the op sees a - # tensor, so this adds no branch to the compiled step. + # Context-only requests already carry the masked end id in the + # store -- FinishReasonsHandler.prepare_for_new_request writes + # the sentinel for them once, at setup -- so the CBA op reads + # it from here with no per-step work. cba_end_ids = self._finish_reasons_handler.store.end_ids_cuda - if beam_search_store.cba is not None and any( - requests[i].is_context_only_request for i in value.indices.tolist() - ): - cba_end_ids = cba_end_ids.clone() - ctx_slots = [ - requests[i].py_seq_slot - for i in value.indices.tolist() - if requests[i].is_context_only_request - ] - cba_end_ids[ - torch.tensor(ctx_slots, device=cba_end_ids.device, dtype=torch.long) - ] = -1 metadata = BeamSearchMetadata( cache_indirection=beam_search_store.cache_indirection, cache_indirection_buffer=beam_search_store.cache_indirection_buffer, diff --git a/tests/unittest/_torch/sampler/test_beam_search.py b/tests/unittest/_torch/sampler/test_beam_search.py index 5b04b8ddbe78..b615a5bf3b5c 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -512,6 +512,107 @@ def test_beam_search_disagg_e2e( ctx_llm.shutdown() +@pytest.mark.threadleak(enabled=False) +def test_beam_search_disagg_first_token_is_end_id( + fixed_params, + input_prompts, + model_kwargs: dict[str, Any], +) -> None: + """A context phase that finishes on its only token still hands off. + + This is the case the handoff is built around, and nothing else reaches it: + an end candidate would normally be pooled and its beam slot refilled, and + the request would be marked finished before it reaches the + disagg-transmission state that builds ContextPhaseParams -- so the + generation server would get no first_gen_tokens at all and never learn the + beam finished. Context-only requests therefore run with their end id + masked to the "no end token" sentinel, which keeps the token in its slot + and keeps the request on the handoff path. + + Triggering it by prompt would be a lottery, and end_id cannot simply be + left unset: these prompts are token ids and the LLM has no tokenizer, so + an unset end_id raises. The end id is picked after the fact instead -- run + once to see what the beams sample, then declare beam 0's token the end id + and rerun, so the context step finishes on its first and only token. + """ + if model_kwargs["sampler_type"] != "TorchSampler": + pytest.skip( + "The context-side end-id mask is a TorchSampler path; the C++ " + "decoder behind TRTLLMSampler pools the end candidate instead.") + + beam_width = fixed_params["max_beam_width"] + base_params = SamplingParams( + max_tokens=fixed_params["max_tokens"], + n=beam_width, + best_of=beam_width, + use_beam_search=True, + end_id=-1, + include_stop_str_in_output=True, + ) + + disagg_kwargs = deepcopy(model_kwargs) + disagg_kwargs |= dict( + disable_overlap_scheduler=True, + cuda_graph_config=None, + kv_cache_config=KvCacheConfig(max_tokens=10000, + enable_block_reuse=True, + enable_partial_reuse=True, + use_kv_cache_manager_v2=True), + cache_transceiver_config=CacheTransceiverConfig( + backend="NIXL", + transceiver_runtime="PYTHON", + kv_transfer_timeout_ms=1000, + kv_transfer_sender_future_timeout_ms=1000, + ), + ) + + prompts = [[1, 2, 3]] + + def _context_first_gen_tokens(llm, params: SamplingParams) -> list[int]: + outputs = llm.generate( + deepcopy(prompts), + sampling_params=deepcopy(params), + disaggregated_params=[ + DisaggregatedParams(request_type="context_only", + disagg_request_id=201) + ], + use_tqdm=False, + ) + assert len(outputs) == len(prompts) + ctx_params = outputs[0].disaggregated_params + assert ctx_params is not None + # None here means the context phase never reached the transmission + # state, so ContextPhaseParams was never built -- the handoff, tokens + # included, was dropped. + assert ctx_params.first_gen_tokens is not None, ( + "context phase produced no first_gen_tokens") + return list(ctx_params.first_gen_tokens) + + ctx_llm = _build_llm(fixed_params, prompts, disagg_kwargs) + try: + with ctx_llm: + # Probe: no end id, so nothing can finish. + baseline_tokens = _context_first_gen_tokens(ctx_llm, base_params) + assert len(baseline_tokens) == beam_width + + end_id = baseline_tokens[0] + end_id_params = deepcopy(base_params) + end_id_params.end_id = end_id + + end_id_tokens = _context_first_gen_tokens(ctx_llm, end_id_params) + + # Beam 0's token is the end id and it is still in the handoff. + # Were it pooled and the slot refilled, some other candidate would + # be here instead; were the request marked finished, there would be + # no handoff to read at all. + assert len(end_id_tokens) == beam_width + assert end_id_tokens[0] == end_id + # Masking is per request, so the other beams are unaffected. + assert end_id_tokens == baseline_tokens + finally: + ctx_llm.shutdown() + + @pytest.mark.parametrize("beam_width", [10]) @pytest.mark.threadleak(enabled=False) def test_beam_search_large_beam_width_regression(