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

tensorrt_llm/_torch/pyexecutor/sampler/sampler.py

@@ -1705,58 +1705,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")
@@ -2357,11 +2360,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()
 
@@ -2461,10 +2470,9 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None:
                         # 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.

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,
 )
@@ -1424,6 +1425,125 @@ 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]):
+            # The beam width must come from the sampling config: it sizes the
+            # request's C++ per-beam state, which set_finished_reason indexes.
+            request = RecordingLlmRequest(
+                request_id=seq_slot,
+                seq_slot=seq_slot,
+                input_tokens=[1],
+                max_new_tokens=10,
+                end_id=2,
+                sampling_config=SamplingConfig(beam_width=beam_width),
+                is_streaming=False,
+            )
+            assert 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):