diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 3a91dd3fcc7b..83adb5eb9a58 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/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. diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index b9dcd85d195a..64401627ca29 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, ) @@ -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):