diff --git a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py index 71dbac467ad3..f742fa22def0 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler.py @@ -1173,6 +1173,12 @@ class SampleStateTensorsHostTorch(SampleStateTensors): finish_reasons: torch.Tensor | None first_finish_reasons: torch.Tensor | None logprobs_state: LogProbsState | None = None + single_step_greedy: bool = False + """Whether `new_tokens` uses the compact `(num_requests,)` layout instead of + `[step, slot, beam]`. Describes these host tensors, so it must live here rather + than on `SampleStateTorch`: under pipeline parallelism only this object crosses + the ring hand-off, and a receiving rank would otherwise pair the compact buffer + with an outer flag left at its default.""" def finish_reasons_list(self) -> FinishReasonsList: """`(num_seq_slots, num_steps)`""" @@ -1187,7 +1193,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 +2363,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 +2751,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..05eebc6a2128 100644 --- a/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py +++ b/tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py @@ -328,7 +328,9 @@ def sample( ) case ("greedy", None): tokens, softmax = greedy_search_sampling_batch(logits, return_probs=return_probs) - temperature = None + # Returns instead of falling through: the other patterns bind + # `temperature` as `float`, so assigning None here does not type check. + return tokens, softmax, None case ( "beam_search", beam_width_in, diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 6f4db28bf2b7..12fdf46312e9 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -171,8 +171,6 @@ full:DGX_B200/perf/test_perf_sanity.py::test_e2e[aggr_upload-gemma4_26b_a4b_nvfp full:DGX_B200/perf/test_perf_sanity.py::test_e2e[aggr_upload-host_perf_llama8b_spec_decode-llama8b_spec_bs1_128_128] SKIP (https://nvbugs/6571408) full:DGX_B200/unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_tinyllama_logits_processor_tp2pp2 SKIP (https://nvbugs/6618096) full:DGX_H100/unittest/llmapi/test_llm_multi_gpu_pytorch.py -m "gpu4" SKIP (https://nvbugs/6618102) -full:DGX_H100/unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_llm_get_stats_pp4[False-False-True] SKIP (https://nvbugs/6618098) -full:DGX_H100/unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_tinyllama_logits_processor_tp2pp2 SKIP (https://nvbugs/6618106) full:GB200/accuracy/test_disaggregated_serving.py::TestLlama3_1_8BInstruct::test_auto_dtype[ctx_block_reuse_only] SKIP (https://nvbugs/6525893) full:GB200/accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite::test_dwdp_accuracy SKIP (https://nvbugs/6276923) full:GB200/accuracy/test_dwdp_disaggregated_serving.py::TestDwdpDeepSeekV3Lite::test_dwdp_accuracy_contention_opt SKIP (https://nvbugs/6276923) 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.