From 6fcfd81176378c4bd8dbd9154c67ffae49372b84 Mon Sep 17 00:00:00 2001 From: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com> Date: Wed, 8 Jul 2026 20:27:42 -0700 Subject: [PATCH] Revert "[None][perf] Move greedy stop checks to host (#15920)" This reverts commit 4f1c3cdcdfe8ebdaa87352823d4fcd708213652f. Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com> --- .../_torch/pyexecutor/sampler/sampler.py | 121 +++++------------- .../_torch/sampler/test_torch_sampler.py | 101 --------------- 2 files changed, 34 insertions(+), 188 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 6726a430ef97..1ad625d6d234 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -1113,8 +1113,6 @@ def finish_reasons_list(self) -> FinishReasonsList: @dataclass(kw_only=True) class SampleStateTorch(SampleState[SampleStateTensorsHostTorch, SampleStateTensors]): beam_history_builders: list[BeamHistoryBuilder | None] | None = None - use_host_stop_criteria: bool = False - """Whether update_requests should evaluate end-ID and length limits on the host.""" @dataclass(kw_only=True, frozen=True) @@ -2454,21 +2452,6 @@ def _can_use_fast_greedy_path(self, requests: list[LlmRequest]) -> bool: return False return True - def _can_use_host_stop_criteria(self, requests: list[LlmRequest]) -> bool: - """Check whether stop criteria can be evaluated from the host token copy. - - The non-speculative, single-beam path already copies sampled tokens to - the host. When no request has stop words, checking end IDs and length - limits on the host avoids the device finish-reason kernels and their - additional D2H copy. - """ - return ( - bool(requests) - and self.max_tokens == 1 - and self.max_beam_width == 1 - and all(not req.py_is_draft and not req.py_stop_words_list for req in requests) - ) - @staticmethod def _meet_max_token_stop_criteria( request: LlmRequest, max_seq_len: int, beam_idx: int = DEFAULT_BEAM_IDX @@ -3736,23 +3719,13 @@ def _maybe_build_beam_history(req_idx: int) -> BeamHistory | None: req.py_rewind_len = 0 else: processed = 1 - if state.use_host_stop_criteria: - new_token = add_token(req, new_tokens_list, beam_idx=DEFAULT_BEAM_IDX) - self._handle_stop_criteria( - req, - new_token, - max_seq_len=self.max_seq_len, - beam_idx=DEFAULT_BEAM_IDX, - ) - num_accepted = 0 - else: - num_accepted = self.process_draft_tokens( - req, - new_tokens_tensor=new_tokens, - new_tokens_list=new_tokens_list, - finish_reasons=finish_reasons, - resource_manager=resource_manager, - ) + num_accepted = self.process_draft_tokens( + req, + new_tokens_tensor=new_tokens, + new_tokens_list=new_tokens_list, + finish_reasons=finish_reasons, + resource_manager=resource_manager, + ) if (actual_draft_len := get_draft_token_length(req)) > 0: req.py_num_accepted_draft_tokens = num_accepted req.py_rewind_len = actual_draft_len - num_accepted @@ -3811,10 +3784,10 @@ def sample_async( ( requests, seq_slots_host, + seq_lens_host, seq_slots_cuda, seq_lens_cuda, new_tokens_host, - use_host_stop_criteria, ) = self._process_requests( scheduled_requests, model_outputs, @@ -3828,8 +3801,7 @@ def sample_async( # Forwarded to _record_sampler_event so SamplerEvent.synchronize # awaits any side-stream D2H copies host-side. side_stream_event: torch.cuda.Event | None = None - if requests and not use_host_stop_criteria: - assert seq_lens_cuda is not None + if requests: beam_search_store = self.store.beam_search_store assert self._use_beam_search == (beam_search_store is not None) # Prepare stop word handling @@ -3895,7 +3867,6 @@ def sample_async( ), sampler_event=sampler_event, beam_history_builders=beam_history_builders, - use_host_stop_criteria=use_host_stop_criteria, ) @staticmethod @@ -4679,12 +4650,7 @@ def _process_requests( new_tokens_cuda: torch.Tensor, num_context_logits_prefix_sum: list[int], ) -> tuple[ - list[LlmRequest], - torch.Tensor, - torch.Tensor, - torch.Tensor | None, - torch.Tensor, - bool, + list[LlmRequest], torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor ]: raw_logits_cuda = model_outputs["logits"] @@ -4697,11 +4663,6 @@ def _process_requests( if return_log_probs: self._prepare_log_probs(sampling_requests) - use_fast_greedy_path = self._can_use_fast_greedy_path(sampling_requests) - use_host_stop_criteria = use_fast_greedy_path and self._can_use_host_stop_criteria( - sampling_requests - ) - seq_slots_host = torch.tensor( [r.py_seq_slot for r in sampling_requests], dtype=torch.int32, @@ -4709,24 +4670,18 @@ def _process_requests( ) # necessary for beam search and max_length checks - seq_lens_host = ( - None - if use_host_stop_criteria - else torch.tensor( - [r.max_beam_num_tokens for r in sampling_requests], - dtype=torch.int32, - pin_memory=prefer_pinned(), - ) + seq_lens_host = torch.tensor( + [r.max_beam_num_tokens for r in sampling_requests], + dtype=torch.int32, + pin_memory=prefer_pinned(), ) - # Cast seq_slots / seq_lens to CUDA exactly once. The fast host stop- - # criteria path only needs seq_slots for scattering sampled tokens; - # all other paths also consume seq_lens in device-side finish handling - # or beam-search metadata. + # Cast seq_slots / seq_lens to CUDA exactly once; consumed by both + # the per-group beam-search metadata builder and the finish-reasons + # handler in sample_async. int64 is required for the index_*_ ops + # downstream. seq_slots_cuda = seq_slots_host.to(device="cuda", dtype=torch.int64, non_blocking=True) - seq_lens_cuda = ( - None if seq_lens_host is None else seq_lens_host.to(device="cuda", non_blocking=True) - ) + seq_lens_cuda = seq_lens_host.to(device="cuda", non_blocking=True) # Handle embedding bias self._apply_embedding_bias( @@ -4741,25 +4696,19 @@ def _process_requests( ) # Fast path for greedy sampling - if use_fast_greedy_path: - if use_host_stop_criteria: - # There is exactly one token and one beam per request, so the - # linearized destination indices are the sequence slots. - batch_dest_indices_cuda = seq_slots_cuda - else: - # Compute destination indices on CPU (same pattern as - # _unbatch_sampling_results). - batch_destination_indexer = _UnpackedStepIndexer( - seq_slots=seq_slots_host, - num_steps=sampling_requests_metadata.req_num_generated_tokens, - steps_dim_size=new_tokens_cuda.size(0), - slots_dim_size=new_tokens_cuda.size(1), - dim_order=_UnpackedStepIndexer.DimOrder.STEP_MAJOR, - index_dtype=torch.int64, - ) - batch_dest_indices_cuda = batch_destination_indexer[:].to( - new_tokens_cuda.device, non_blocking=True - ) + if self._can_use_fast_greedy_path(sampling_requests): + # Compute destination indices on CPU (same pattern as _unbatch_sampling_results) + batch_destination_indexer = _UnpackedStepIndexer( + seq_slots=seq_slots_host, + num_steps=sampling_requests_metadata.req_num_generated_tokens, + steps_dim_size=new_tokens_cuda.size(0), + slots_dim_size=new_tokens_cuda.size(1), + dim_order=_UnpackedStepIndexer.DimOrder.STEP_MAJOR, + index_dtype=torch.int64, + ) + batch_dest_indices_cuda = batch_destination_indexer[:].to( + new_tokens_cuda.device, non_blocking=True + ) # Get d2t tensor if present d2t = model_outputs.get("d2t", None) @@ -4777,10 +4726,10 @@ def _process_requests( return ( sampling_requests, seq_slots_host, + seq_lens_host, seq_slots_cuda, seq_lens_cuda, new_tokens_host, - use_host_stop_criteria, ) # Indexer for accessing tokens in 'logits_cuda', corresponding to the @@ -4793,8 +4742,6 @@ def _process_requests( ) # Perform sampling in batches - assert seq_lens_host is not None - assert seq_lens_cuda is not None batched_sampling_result = self._sample_batched_by_strategy( logits_cuda, sampling_requests, @@ -4832,10 +4779,10 @@ def _process_requests( return ( sampling_requests, seq_slots_host, + seq_lens_host, seq_slots_cuda, seq_lens_cuda, new_tokens_host, - use_host_stop_criteria, ) @override diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index b76d4aba0d16..87e6b4b8dc01 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -866,107 +866,6 @@ def test_write_finish_reasons(cls): run_test_with_warmup(uut_provider, max_sync_s=0.5) - @pytest.mark.parametrize( - "sampled_token,end_id,max_new_tokens,expected_reason", - [ - pytest.param(42, 99, 5, None, id="not-finished"), - pytest.param(99, 99, 5, FinishReason.END_ID, id="end-id"), - pytest.param(42, 99, 1, FinishReason.LENGTH, id="length"), - pytest.param( - 99, - 99, - 1, - FinishReason.END_ID, - id="end-id-precedes-length", - ), - ], - ) - def test_host_stop_criteria_fast_path( - self, - mocker, - sampled_token: int, - end_id: int, - max_new_tokens: int, - expected_reason: FinishReason | None, - ): - sampler = TorchSampler( - TorchSampler.Args( - max_seq_len=20, - max_draft_len=0, - max_total_draft_tokens=0, - max_num_sequences=1, - max_beam_width=1, - ) - ) - request = LlmRequest( - request_id=0, - seq_slot=0, - input_tokens=[1, 2], - max_new_tokens=max_new_tokens, - end_id=end_id, - sampling_config=SamplingConfig(), - is_streaming=False, - ) - - setup_requests = ScheduledRequests() - setup_requests.context_requests_last_chunk = [request] - sampler.setup_sampler_step(setup_requests) - - scheduled_requests = ScheduledRequests() - scheduled_requests.generation_requests = [request] - logits = torch.full((1, 128), -1.0, dtype=torch.float32, device="cuda") - logits[0, sampled_token] = 1.0 - - finish_by = mocker.spy(request, "finish_by") - write_finish_reasons = mocker.patch.object( - sampler._finish_reasons_handler, - "write_finish_reasons", - side_effect=AssertionError("device finish reasons should be skipped"), - ) - - state = sampler.sample_async( - scheduled_requests, - model_outputs={"logits": logits}, - num_context_logits_prefix_sum=[0], - ) - assert state.use_host_stop_criteria - assert state.host is not None - assert state.host.finish_reasons is None - - sampler.update_requests(state) - - write_finish_reasons.assert_not_called() - assert request.get_tokens(0)[-1] == sampled_token - if expected_reason is None: - finish_by.assert_not_called() - else: - finish_by.assert_called_once_with(expected_reason, 0) - - @pytest.mark.parametrize( - "sample_request", - [ - pytest.param( - SimpleNamespace(py_is_draft=False, py_stop_words_list=[[42], [1]]), - id="stop-words", - ), - pytest.param( - SimpleNamespace(py_is_draft=True, py_stop_words_list=None), - id="draft-request", - ), - ], - ) - def test_host_stop_criteria_fast_path_fallback(self, sample_request): - sampler = TorchSampler( - TorchSampler.Args( - max_seq_len=20, - max_draft_len=0, - max_total_draft_tokens=0, - max_num_sequences=1, - max_beam_width=1, - ) - ) - assert not sampler._can_use_host_stop_criteria([cast(LlmRequest, sample_request)]) - @classmethod def test_are_stop_words_isnt_called_when_no_stop_words(cls, monkeypatch: pytest.MonkeyPatch): """We don't want to call are_stop_words when there are no stop words because it's expensive"""