From 0d6d0e9112bd7f277b8a2d8bcc0b8034ce022b69 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Mon, 17 Aug 2026 03:50:24 -0700 Subject: [PATCH] [nvbugs/6618100][fix] Carry the stable-greedy layout marker across the PP hand-off The stable-greedy fast path in _process_requests returns a compact 1-D new_tokens_host of shape (num_requests,) rather than the usual [step, slot, beam], and flagged that layout only on the outer SampleStateTorch. Pipeline parallelism transfers just sample_state.host around the ring, so non-final ranks rebuilt the state object with single_step_greedy defaulted to False and indexed the flat buffer as 3-D, raising "TypeError: 'int' object is not subscriptable" once per non-final PP stage and aborting the executor. Move the marker onto SampleStateTensorsHostTorch -- the object actually transferred -- so the buffer is self-describing and send/recv stay symmetric with no wire-format change, and dispatch on either copy. Also return the greedy branch's result directly in sampler_strategy.sample() instead of assigning temperature = None, which mypy rejects as incompatible with the float bound by the surrounding capture patterns. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/sampler/sampler.py | 10 +++++++--- .../_torch/pyexecutor/sampler/sampler_strategy.py | 2 +- tests/unittest/_torch/sampler/test_torch_sampler.py | 4 ++-- 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 71dbac467ad3..2c4c2ce6ff00 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -1173,6 +1173,11 @@ class SampleStateTensorsHostTorch(SampleStateTensors): finish_reasons: torch.Tensor | None first_finish_reasons: torch.Tensor | None logprobs_state: LogProbsState | None = None + single_step_greedy: bool = False + """Marks `new_tokens` as the compact 1-D `(num_requests,)` layout instead of + `[step, slot, beam]`. Must stay on this dataclass, the only part of the sample + state pipeline parallelism transfers (`_ring_broadcast_sample_state` sends just + `host`), so a receiving rank cannot pair the buffer with a stale default.""" def finish_reasons_list(self) -> FinishReasonsList: """`(num_seq_slots, num_steps)`""" @@ -1187,7 +1192,6 @@ def finish_reasons_list(self) -> FinishReasonsList: @dataclass(kw_only=True) class SampleStateTorch(SampleState[SampleStateTensorsHostTorch, SampleStateTensors]): beam_history_builders: list[BeamHistoryBuilder | None] | None = None - single_step_greedy: bool = False class _SideStreamCopier: @@ -2358,7 +2362,7 @@ def update_requests( assert state.host is not None # Reuse sample_async's qualification instead of rechecking every # request after the asynchronous sample completes. - if state.single_step_greedy: + if state.host.single_step_greedy: self._update_requests_single_beam_single_step(state) return @@ -2746,10 +2750,10 @@ def sample_async( finish_reasons=finish_reasons_host, first_finish_reasons=first_finish_reasons_host, logprobs_state=logprobs_state, + single_step_greedy=single_step_greedy, ), sampler_event=sampler_event, beam_history_builders=beam_history_builders, - single_step_greedy=single_step_greedy, ) @staticmethod diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py index 51b69c1d4c24..5bd322205a6d 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py @@ -328,7 +328,7 @@ def sample( ) case ("greedy", None): tokens, softmax = greedy_search_sampling_batch(logits, return_probs=return_probs) - temperature = None + return tokens, softmax, None case ( "beam_search", beam_width_in, diff --git a/tests/unittest/_torch/sampler/test_torch_sampler.py b/tests/unittest/_torch/sampler/test_torch_sampler.py index 64401627ca29..978f9934422e 100644 --- a/tests/unittest/_torch/sampler/test_torch_sampler.py +++ b/tests/unittest/_torch/sampler/test_torch_sampler.py @@ -961,8 +961,8 @@ def test_single_step_greedy_updates_finish_reasons_and_filters_completed_request new_tokens=new_tokens, finish_reasons=None, first_finish_reasons=None, + single_step_greedy=True, ), - single_step_greedy=True, ) sampler.update_requests(state) @@ -1986,7 +1986,7 @@ def _mock_filter(self, requests: ScheduledRequests) -> list[LlmRequest]: sample_state.sampler_event.synchronize() assert sample_state.host is not None host_new_tokens = sample_state.host.new_tokens - if sample_state.single_step_greedy: + if sample_state.host.single_step_greedy: # The stable greedy path copies one token per active request instead of # the full [step, slot, beam] buffer. This fixture uses dense sequence # slots, so restore that layout before comparing sampling results.