[https://nvbugs/6618106][fix] Carry single_step_greedy on SampleStateTensorsHostTorch

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

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,

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)

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.