[TRTLLM-13234][feat] Complete TorchSampler beam search: length_penalty, diversity_rate, early_stopping, VBWS, and CBA performance

cpp/tensorrt_llm/batch_manager/llmRequest.cpp

@@ -37,13 +37,16 @@ runtime::SizeType32 GenericLlmRequest<TTensor, TStream>::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<int>(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<int>(requestBeamWidthArray.size())) - 1, 0);
+        beamWidth = requestBeamWidthArray[index];
     }
     return beamWidth;
 }

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.

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(

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

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)