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; } diff --git a/docs/source/features/sampling.md b/docs/source/features/sampling.md index 4f8eb101f450..277aa64eeddc 100644 --- a/docs/source/features/sampling.md +++ b/docs/source/features/sampling.md @@ -274,6 +274,38 @@ 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. 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, 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: @@ -291,6 +323,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 c68b83c2abde..9d9a744e1e94 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -868,6 +868,38 @@ 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. + + 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. + + 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: + 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]) + 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..3715962c18a4 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. 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_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). @@ -7832,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) @@ -7847,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/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 62b7c6488ee9..9dd6c2bbeb4a 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -4974,11 +4974,57 @@ 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!") + + # 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.") + + # 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.") # Check token ID ranges self._validate_token_id_range(request) @@ -6355,7 +6401,16 @@ 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. + + 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 @@ -6430,6 +6485,31 @@ 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, 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. + + # 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 @@ -6880,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/beam_search.py b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py new file mode 100644 index 000000000000..65e733b62d99 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/sampler/beam_search.py @@ -0,0 +1,1488 @@ +# 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) +beam-search logic for TorchSampler, split out of ``ops.vanilla``. 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, Any, Callable, NamedTuple, Optional, TypeAlias, cast + +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 + +from ..llm_request import LlmRequest, LlmRequestState +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 + +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.""" + + 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``. + + 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: + 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). + + 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 IS_FLASHINFER_AVAILABLE and 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. + + 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 + (``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 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 + 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 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 + """[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.""" + 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.""" + 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_``.""" + 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 beam-search 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), + 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), + ) + + def ensure_cba(self) -> _CBAFields: + """Allocate the CBA tensors on first use and return them. + + 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 + 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.""" + + +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_cba.""" + + 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 + pending_harvest: torch.Tensor + predecessor_beams: torch.Tensor + beam_idx_arange: 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). + 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 for every beam-search request + regardless of its early_stopping mode; None until the first one.""" + + +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, + row_stride: int | None = None, + 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. + + 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`` + 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]" + 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, 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 + 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 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 + 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 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. + + 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 + + +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, + pending_harvest: 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 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") + + 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) + 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. + # 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) + + # 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) + # `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.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)) + # 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). + 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). +# +# 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( + logits: torch.Tensor, + *, + beam_width_in: int, + beam_width_out: int, + row_stride: int | None = None, + 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). 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 + 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. + + 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 + assert cba is not None, "CBA metadata is required for beam search" + 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, + row_stride=row_stride, + 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.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] + + # --- 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, + pending_harvest=args.pending_harvest, + 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 + # 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. + # + # 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) & (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: + stop_window[:, slots, :num_beams] = reordered_window + + # --- 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] + 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). + # 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] = beam_log_probs + 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.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) + # 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: + 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 _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 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() + ) + + +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], + 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._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). + """ + # Every beam-search request runs on this path. + cba_requests = list(requests) + 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 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] + 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_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. + """ + # 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() + ) + 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 = [_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 + + 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/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py b/tensorrt_llm/_torch/pyexecutor/sampler/finish_reasons.py index e79281176391..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 @@ -528,6 +544,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 +584,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 +656,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 +731,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/ops/flashinfer.py b/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py index 6ee643336eb7..ae43820dd1df 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/ops/flashinfer.py @@ -64,6 +64,28 @@ 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. + """ + topk_values, topk_indices = flashinfer.top_k(values, k, sorted=True, deterministic=True) + return topk_values, topk_indices + + @_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..fcf6e768a8dc 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,12 @@ 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 .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 @@ -101,6 +100,7 @@ DEFAULT_BEAM_IDX, DEFAULT_STEP_IDX, FinishReasonsList, + _get_beam_width_out, _get_max_beam_width, _request_get_sampling_params, add_token, @@ -108,9 +108,11 @@ request_random_seed, ) from .sampler_strategy import ( - BEAM_SEARCH_PAD_TOKEN, GREEDY, + BeamHistory, BeamSearchMetadata, + BeamSearchStore, + CBAState, FlashInferGroupedStrategySampler, Fusions, GenericStrategyKeyType, @@ -156,6 +158,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 +179,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 +1150,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 +1190,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 +1391,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 +1398,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 +1427,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 +1561,23 @@ 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, + 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. @@ -1800,58 +1704,61 @@ def _handle_stop_criteria( return False - def _handle_finish_reasons_impl( - self, - request: LlmRequest, - beam_width: int, - finish_reasons: torch.Tensor, - finish_reasons_list: list[int], - ) -> bool: - """Check if all beams of a request have finished and set the request state accordingly + @staticmethod + def _finished_beam_prefix_lengths(finish_reasons: torch.Tensor) -> list[int]: + """Count the leading finished beams of every slot in one batched reduction. + + A request is complete once all of the beams it actually uses have a finish + reason, i.e. once its leading ``py_beam_width`` entries are all set. Rather + than reducing each request's row separately, reduce the whole tensor once + and return, per slot, how many leading beams have finished. The per-request + check then degenerates to ``prefix_length >= beam_width``, which needs no + tensor work and stays correct for mixed beam widths regardless of what the + columns past a request's width hold. Args: - request: LlmRequest. The request to check. - beam_width: int. The beam width of the request. - finish_reasons: torch.Tensor. Shape: (beam_width) - The finish reasons for each beam. - finish_reasons_list: list[int]. The finish reasons for each beam. + finish_reasons: Shape ``(max_batch_size, max_beam_width)``. The finish + reasons of every beam of every slot. + Returns: - True if all beams have finished, False otherwise. + Per slot, the number of leading beams whose finish reason is set. """ - if (finish_reasons[:beam_width] != FinishReason.NOT_FINISHED.value).sum() == beam_width: - request.state = LlmRequestState.GENERATION_COMPLETE - for beam_idx in range(beam_width): - request.set_finished_reason( - FinishReason(finish_reasons_list[beam_idx]), - beam_idx, - ) - return True - return False + unfinished = finish_reasons == FinishReason.NOT_FINISHED.value + # A beam belongs to the finished prefix iff no unfinished beam precedes it + # and it is finished itself, i.e. iff the running count of unfinished beams + # up to and including it is still zero. Counting those positions yields the + # prefix length directly, and needs no special case for a fully finished + # row (every position counts) or a row finishing at beam 0 (none do). + return (unfinished.cumsum(dim=1) == 0).sum(dim=1).tolist() def _handle_first_finish_reasons( self, request: LlmRequest, - finish_reasons: torch.Tensor, + finished_beam_prefix_lengths: list[int], finish_reasons_list: list[list[int]], ) -> bool: """Check if all beams of a request have finished and set the request state accordingly Args: request: LlmRequest. The request to check. - finish_reasons: torch.Tensor. Shape: (max_batch_size, max_beam_width) - The finish reasons for each beam. + finished_beam_prefix_lengths: Per slot, the number of leading beams that + have finished, as returned by ``_finished_beam_prefix_lengths``. finish_reasons_list: list[list[int]]. The finish reasons for each beam. Returns: True if all beams have finished, False otherwise. """ assert request.py_seq_slot is not None beam_width = request.py_beam_width - return self._handle_finish_reasons_impl( - request, - beam_width, - finish_reasons[request.py_seq_slot, :beam_width], - finish_reasons_list[request.py_seq_slot], - ) + if finished_beam_prefix_lengths[request.py_seq_slot] < beam_width: + return False + request.state = LlmRequestState.GENERATION_COMPLETE + request_finish_reasons = finish_reasons_list[request.py_seq_slot] + for beam_idx in range(beam_width): + request.set_finished_reason( + FinishReason(request_finish_reasons[beam_idx]), + beam_idx, + ) + return True @staticmethod @nvtx_range("update_original_tokens") @@ -2078,6 +1985,9 @@ def validate_request(self, request: LlmRequest) -> None: raise ValueError( "Beam search only supports returning the sampled logprob per token" ) + # 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") @@ -2113,15 +2023,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 +2043,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 +2068,18 @@ 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 beam-search request: every + # early_stopping mode runs on that path. + 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 +2211,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], @@ -2611,6 +2250,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 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 metadata = BeamSearchMetadata( cache_indirection=beam_search_store.cache_indirection, cache_indirection_buffer=beam_search_store.cache_indirection_buffer, @@ -2619,9 +2263,35 @@ 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, - seq_offsets=beam_search_store.seq_offsets, beam_idx_arange=beam_search_store.beam_idx_arange, + 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. + cba=None + if beam_search_store.cba is None + else CBAState( + 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, + 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 +2313,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 +2329,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() @@ -2764,11 +2358,17 @@ def update_requests( new_tokens = state.host.new_tokens finish_reasons = state.host.finish_reasons_list() - first_finish_reasons = ( - state.host.first_finish_reasons.tolist() - if state.host.first_finish_reasons is not None - else [] - ) + first_finish_reasons_host = state.host.first_finish_reasons + if first_finish_reasons_host is not None: + first_finish_reasons = first_finish_reasons_host.tolist() + # Reduce every slot at once; the per-request loop below only reads the + # result and updates the request objects. + finished_beam_prefix_lengths = self._finished_beam_prefix_lengths( + first_finish_reasons_host + ) + else: + first_finish_reasons = [] + finished_beam_prefix_lengths = [] new_tokens_list = new_tokens.tolist() @@ -2838,23 +2438,45 @@ 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: - self._finalize_beam(req, beam_history) + # 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: - 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) - first_finish_reasons_host = state.host.first_finish_reasons assert first_finish_reasons_host is not None self._handle_first_finish_reasons( - req, first_finish_reasons_host, first_finish_reasons + req, finished_beam_prefix_lengths, first_finish_reasons ) 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 @@ -3066,6 +2688,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) @@ -3077,7 +2702,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 ) @@ -3772,9 +3397,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() @@ -4026,7 +3661,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/tensorrt_llm/_torch/pyexecutor/sampler/sampler_common.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_common.py index f7ba87258654..e787d69f970d 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] @@ -63,9 +76,15 @@ 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 + 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: @@ -153,7 +172,15 @@ 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( + 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, @@ -162,10 +189,14 @@ 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, 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..fc04afbb51e4 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py @@ -24,14 +24,25 @@ 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_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, @@ -46,10 +57,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 +79,16 @@ # 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 +116,30 @@ # (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 + # 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) 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 +238,16 @@ 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), + row_stride=params.row_stride or params.beam_width_in, + ) # NB: not greedy, hence top_p != 0 if specified top_p = top_p or 1.0 @@ -290,16 +329,34 @@ 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, + *_, + ): + 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" ) - tokens, softmax = beam_search_sampling_batch( + # 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) @@ -953,25 +1010,117 @@ 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 + ``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. + """ + + @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 + row_stride: 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, + row_stride: 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._row_stride = row_stride 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( + # 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 ) - 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, + row_stride=row_stride, + 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.row_stride, + fields.temperature, + fields.length_penalty, + fields.diversity_rate, + ) @override def sample( @@ -984,23 +1133,87 @@ 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 beam_search_sampling_batch( + 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 CBABeamSearchStep(BeamSearchStep): + """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, + beam_width_in: int, + beam_width_out: int, + row_stride: 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, + row_stride, + 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.row_stride, + 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, + row_stride=self._row_stride, 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(), ) - class BeamSearchWithProbs(BeamSearchMixin, StrategyImplWithProbs): - pass - - class BeamSearchSampleOnly(BeamSearchMixin, StrategyImplSampleOnly): - pass - _STRATEGY_KEY_TYPE: TypeAlias = ( Literal["temperature"] @@ -1009,7 +1222,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 +1243,21 @@ 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)) + # 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), + ) case _: raise NotImplementedError("Unsupported strategy encountered") @@ -1040,7 +1266,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 +1308,12 @@ 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, _, _): beam_width_in = beam_width_in_key - strategy_impl_cls = _StrategyImpls.BeamSearchWithProbs + # 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: @@ -1101,16 +1330,26 @@ 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, _, _): beam_width_in = beam_width_in_key - strategy_impl_cls = _StrategyImpls.BeamSearchSampleOnly + strategy_impl_cls = _StrategyImpls.CBABeamSearchStep 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) + # 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..9d9ecb4d2864 100644 --- a/tensorrt_llm/sampling_params.py +++ b/tensorrt_llm/sampling_params.py @@ -234,16 +234,16 @@ 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. 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/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 0cce73d0b5f8..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): @@ -530,8 +549,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[EarlyStopping] = None stop_token_ids: Optional[List[int]] = Field(default_factory=list) include_stop_str_in_output: bool = False ignore_eos: bool = False @@ -616,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, @@ -886,8 +909,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[EarlyStopping] = None stop_token_ids: Optional[List[int]] = Field(default_factory=list) include_stop_str_in_output: bool = False ignore_eos: bool = False @@ -1032,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/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 82f879f2debb..b615a5bf3b5c 100644 --- a/tests/unittest/_torch/sampler/test_beam_search.py +++ b/tests/unittest/_torch/sampler/test_beam_search.py @@ -12,13 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +import dataclasses +import functools import gc import os import pathlib as _pl +import types from contextlib import contextmanager, nullcontext from copy import deepcopy -from dataclasses import dataclass -from types import SimpleNamespace from typing import Any, Callable, Generator, cast import pytest @@ -32,18 +33,22 @@ 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.py_executor import PyExecutor from tensorrt_llm._torch.pyexecutor.sampler import (BeamHistory, SampleStateTorch, TorchSampler) -from tensorrt_llm._torch.pyexecutor.sampler.logprobs import \ - convert_logprobs_tensor_to_list +from tensorrt_llm._torch.pyexecutor.sampler.beam_search import ( + CBAGroupHost, _gather_beam_path, _prepare_beam_history_cba, finalize_beam) 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_cba) 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) +from tensorrt_llm.executor.result import CompletionOutput, GenerationResult from tensorrt_llm.llmapi import (CacheTransceiverConfig, CudaGraphConfig, KvCacheConfig) @@ -161,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[ @@ -175,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)}" @@ -201,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 @@ -245,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}" @@ -287,31 +318,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, @@ -377,56 +383,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]) @@ -492,10 +448,15 @@ 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 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"], n=fixed_params["max_beam_width"], @@ -503,7 +464,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) @@ -522,25 +482,135 @@ def test_beam_search_disagg_e2e( ), ) - partial_reuse_prompts = [[1, 2, 3], [1, 5, 6]] + prompts = [[1, 2, 3]] + ctx_llm = _build_llm(fixed_params, prompts, disagg_kwargs) + try: + with ctx_llm: + # 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 + 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() - ctx_llm = _build_llm(fixed_params, partial_reuse_prompts, disagg_kwargs) - gen_llm = _build_llm(fixed_params, partial_reuse_prompts, disagg_kwargs) + +@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, 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) + 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() - gen_llm.shutdown() @pytest.mark.parametrize("beam_width", [10]) @@ -613,9 +683,230 @@ 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(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 -> + 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 + 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 + # its last entry for most of the run: that clamp is the interesting part, + # since getting it wrong reads past the array. + # + # 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(), + 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) + + # 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( + 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, + **extra_params, + 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}") + + # 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))}") + + # 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 ########################################################################### + + +@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 + 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: Any, **kwargs: Any) -> Any: + with torch.device("cuda"): + return fn(*args, **kwargs) + + # 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: # Test Parameters for the update_beam_history and finish_beams tests beam_width = 3 @@ -633,210 +924,168 @@ class GeneralTestParams: vocab_size = 100 -def test_beam_search_sampling_batch_basic(): - """Test basic beam search sampling functionality.""" +_CBA_NEG_INF = float("-inf") - 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, - ) +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. - # 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, - ) - - # 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 were updated: TODO -- This test currently always passes, as finished beams is always 0. - 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]) - # 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)) - - -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) + 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("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), + ) - def make_metadata(seq_len: int) -> BeamSearchMetadata: + # 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, prompt_len + 2), - dtype=torch.int32), + (max_batch_size, beam_width, seq_len + 1), dtype=torch.int32), cache_indirection_buffer=torch.full( - (max_batch_size, beam_width, prompt_len + 2), + (max_batch_size, beam_width, seq_len + 1), -1, dtype=torch.int32), cum_log_probs=torch.zeros((max_batch_size, beam_width), @@ -845,139 +1094,601 @@ def make_metadata(seq_len: int) -> 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), dtype=torch.int32), - seq_offsets=seq_offsets, - beam_idx_arange=beam_idx_arange, + beam_idx_arange=torch.arange(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, + 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 = _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_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 + + # 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] + + # 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_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) + + +# 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.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.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), + # 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( + 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.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.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.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.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, + ), + ) + 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, ) - 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, + + # 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, ) - 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, - ))))) - 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, - ), + assert m.cba.batch_dones[0].item() is expect_done + # 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). + assert torch.allclose(m.cba.cba_normed_scores[0], torch.tensor([-0.1, + -1.5])) + + +@_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 + # 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) + 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, + beam_width_out=K, + beam_search_args=m, + temperature=1.0, + early_stopping=0, + length_penalty=penalty, + return_probs=False, ) - 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, + _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 + # 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}") + + +@_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 + 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) + 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] - 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, + +@_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, ) - assert not torch.allclose( - unseeded_metadata.cum_log_probs[seq_slots, :beam_width], - continuous_cum_log_probs) + # 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 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] def create_default_request(test_params: GeneralTestParams) -> LlmRequest: @@ -1028,163 +1739,341 @@ def create_default_sampler(test_params: GeneralTestParams) -> TorchSampler: return sampler -def test_create_beam_history(): - """Test TorchSampler._create_beam_history method. +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) + - This test verifies that beam history is correctly reconstructed by following - the cache_indirection backwards to obtain the correct token sequence. +@pytest.mark.parametrize( + "beam_width_array, expected", + [ + # Index is (iteration - 1), clamped at both ends. + ([2, 3, 4], [2, 2, 3, 4]), + ([2, 2, 4], [2, 2, 2, 4]), + ], + 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. + + 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 = [] + 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. + + 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) + # 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_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) + + 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 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 CppLlmRequest.get_beam_width_by_iter( + 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. """ - @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) + 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 + ) - # 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 - - # 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() + if accepted: + validate(request) + else: + with pytest.raises(ValueError, match="decreases"): + validate(request) - # test - def _uut(res=res): - res.result = UutResult( - beam_history_builder=sampler._prepare_beam_history( - request, - finish_reasons=torch.ones((beam_width, ), dtype=torch.int), - d2h_copier=sampler._copy_to_host, - ), ) - yield _uut +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 - 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")) + 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) - run_test_with_warmup(_uut_provider, max_sync_s=1) + # 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 + 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_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_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. + + 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. + """ + + test_params = GeneralTestParams() + beam_width = test_params.beam_width + num_generated_tokens = test_params.num_generated_tokens + + 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) + + +def test_cba_finalize_merges_pool_and_orders_by_score(): + """CBA finalization ranks pool entries against the live beams. + + _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) + + # 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, + ) + + 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(): @@ -1250,7 +2139,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 +2174,95 @@ 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, 2, torch.ones(1), None, None) + + @staticmethod + @pytest.mark.parametrize( + "early_stopping", + [ + BeamSearchEarlyStop.TRUE, + BeamSearchEarlyStop.FALSE, + BeamSearchEarlyStop.NEVER, + ], + ) + @_kernel_test + 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")) + 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 +2318,76 @@ 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, - )) + sampling_params=SamplingParams(**params)) 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): + @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, + ): + # 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: 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, - )) + 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) @@ -1404,17 +2402,33 @@ 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=".*is not equal to 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..fd3a38b53257 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. @@ -30,19 +30,21 @@ """ 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 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 +159,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 +171,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( @@ -337,6 +337,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. @@ -369,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], @@ -379,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. @@ -392,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) diff --git a/tests/unittest/_torch/sampler/test_logits_logprobs.py b/tests/unittest/_torch/sampler/test_logits_logprobs.py index dc9bc4808d25..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 @@ -1362,7 +1378,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. diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index 4a29f602bba6..5d0bd7a1936c 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -38,6 +38,7 @@ from tensorrt_llm._torch.pyexecutor.llm_request import ( LlmRequest, + LlmRequestState, convert_wordlist, get_draft_token_length, ) @@ -146,6 +147,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 +173,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 +566,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 @@ -1418,6 +1425,123 @@ def setup_sampler_step_with_size_check(self, scheduled_requests: ScheduledReques ) run_test_with_warmup(uut_provider_with_resize_on_demand, max_sync_s=None) + @staticmethod + def _all_beams_finished_reference(row: torch.Tensor, beam_width: int) -> bool: + """The per-request reduction the batched prefix count replaces.""" + return bool( + (row[:beam_width] != FinishReason.NOT_FINISHED.value).sum().item() == beam_width + ) + + def test_finished_beam_prefix_lengths_matches_per_request_reduction(self): + """The batched prefix count answers the per-request question for every width.""" + store_width = 4 + reasons = [ + FinishReason.NOT_FINISHED.value, + FinishReason.END_ID.value, + FinishReason.STOP_WORDS.value, + FinishReason.LENGTH.value, + ] + rows = list(product(reasons, repeat=store_width)) + finish_reasons = torch.tensor(rows, dtype=torch.int32) + + prefix_lengths = TorchSampler._finished_beam_prefix_lengths(finish_reasons) + + assert len(prefix_lengths) == len(rows) + for row, prefix_length in zip(finish_reasons, prefix_lengths): + for beam_width in range(1, store_width + 1): + assert (prefix_length >= beam_width) == self._all_beams_finished_reference( + row, beam_width + ), f"row={row.tolist()} beam_width={beam_width}" + + def test_finished_beam_prefix_lengths_ignores_columns_past_beam_width(self): + """Reasons beyond a request's beam width must not complete it, or vice versa.""" + # Slot 0 uses 2 beams and both finished; the padding columns are unfinished. + # Slot 1 uses 2 beams, only the second finished; the padding columns are set. + finish_reasons = torch.tensor( + [ + [FinishReason.END_ID.value, FinishReason.LENGTH.value, 0, 0], + [ + 0, + FinishReason.END_ID.value, + FinishReason.END_ID.value, + FinishReason.END_ID.value, + ], + ], + dtype=torch.int32, + ) + + prefix_lengths = TorchSampler._finished_beam_prefix_lengths(finish_reasons) + + assert prefix_lengths[0] >= 2 + assert prefix_lengths[1] < 2 + + def test_handle_first_finish_reasons_completes_only_fully_finished_requests(self): + """Requests are completed, and their per-beam reasons recorded, only when all + of their own beams finished -- across differing beam widths in one batch.""" + sampler = object.__new__(TorchSampler) + store_width = 4 + # Slot 0: beam_width 2, both finished -> completes. + # Slot 1: beam_width 4, first beam unfinished -> stays running. + # Slot 2: beam_width 1, finished -> completes. + finish_reasons = torch.tensor( + [ + [FinishReason.END_ID.value, FinishReason.LENGTH.value, 0, 0], + [ + 0, + FinishReason.END_ID.value, + FinishReason.END_ID.value, + FinishReason.END_ID.value, + ], + [FinishReason.STOP_WORDS.value, 0, 0, 0], + ], + dtype=torch.int32, + ) + assert finish_reasons.size(1) == store_width + prefix_lengths = TorchSampler._finished_beam_prefix_lengths(finish_reasons) + finish_reasons_list = finish_reasons.tolist() + + class RecordingLlmRequest(LlmRequest): + """LlmRequest that records the per-beam reasons the sampler sets.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.recorded_reasons: list[tuple[int, FinishReason]] = [] + + def set_finished_reason(self, finish_reason: FinishReason, beam: int) -> None: + self.recorded_reasons.append((beam, finish_reason)) + super().set_finished_reason(finish_reason, beam) + + requests = [] + for seq_slot, beam_width in enumerate([2, 4, 1]): + request = RecordingLlmRequest( + request_id=seq_slot, + seq_slot=seq_slot, + input_tokens=[1], + max_new_tokens=10, + end_id=2, + sampling_config=SamplingConfig(), + is_streaming=False, + ) + request.py_beam_width = beam_width + requests.append(request) + + completed = [ + sampler._handle_first_finish_reasons(request, prefix_lengths, finish_reasons_list) + for request in requests + ] + + assert completed == [True, False, True] + assert requests[0].state == LlmRequestState.GENERATION_COMPLETE + assert requests[1].state != LlmRequestState.GENERATION_COMPLETE + assert requests[2].state == LlmRequestState.GENERATION_COMPLETE + # Only the request's own beams are reported, in beam order. + assert requests[0].recorded_reasons == [ + (0, FinishReason.END_ID), + (1, FinishReason.LENGTH), + ] + assert requests[1].recorded_reasons == [] + assert requests[2].recorded_reasons == [(0, FinishReason.STOP_WORDS)] + @pytest.mark.parametrize("min_p", [0.0, 0.1, 0.5, 0.9]) def test_min_p_renorm_probs(min_p: float): @@ -3455,6 +3579,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) diff --git a/tests/unittest/api_stability/references/trtllm_serve_api.yaml b/tests/unittest/api_stability/references/trtllm_serve_api.yaml index 99d263914f05..34f7e05bd495 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[Union[bool, Literal['never']]] + 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[Union[bool, Literal['never']]] + default: null status: stable required: false stop_token_ids: