From f49d5b737a5704ba30fae2c919124d9195136261 Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Mon, 17 Aug 2026 04:05:00 -0700 Subject: [PATCH 1/4] [nvbugs/6618106][fix] Carry the stable-greedy layout flag on the PP-transmitted host tensors The stable-greedy fast path in `_process_requests` returns a compact 1-D `new_tokens_host` of shape `(num_requests,)` instead of the usual `[step, slot, beam]` buffer, but flagged that layout only on the outer `SampleStateTorch`. Under pipeline parallelism, however, `_ring_broadcast_sample_state` transmits just `sample_state.host`, and `_forward_step_inter_pp` rebuilds the state object locally without the flag, so a non-final PP rank paired the received flat buffer with the dataclass default `single_step_greedy=False`. It then fell into the batched fast path and indexed the flat token list as 3-D: TypeError: 'int' object is not subscriptable printed once per non-final PP stage. The reported "Test leaked" failure is a downstream artifact of the aborted run, not the defect. Move the marker onto `SampleStateTensorsHostTorch` -- the object actually pickled across the ring -- and dispatch on either copy in `update_requests`. This keeps send/recv symmetric without changing the wire-format tuple on both sides, and needs no special-casing for sampler types lacking the field. In sampler_strategy.py, return `(tokens, softmax, None)` directly from the greedy branch instead of assigning `temperature = None` and falling through. The tuple patterns bind `temperature` as `float`, so that assignment is a type error reported by the `type-check` hook, which is `pass_filenames: false` and therefore lints this file for any edit under the sampler directory. Annotating `temperature` as `float | None` would clear the local run but make the `case` capture patterns ill-typed under the pinned mypy 1.19.1 that CI installs; the early return is clean under both 1.19.1 and 1.20.1 and is a no-op at runtime, since the fall-through already returned `cast(float, None)`. Also drops this bug's waives.txt line. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/sampler/sampler.py | 11 ++++++++--- .../_torch/pyexecutor/sampler/sampler_strategy.py | 4 +++- tests/integration/test_lists/waives.txt | 1 - tests/unittest/_torch/sampler/test_torch_sampler.py | 4 ++-- 4 files changed, 13 insertions(+), 7 deletions(-) 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..8f54d011ddaf 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -172,7 +172,6 @@ full:DGX_B200/perf/test_perf_sanity.py::test_e2e[aggr_upload-host_perf_llama8b_s 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. From ed860b2af07817c70f6fbb521cfab6ac7e6f5139 Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Mon, 17 Aug 2026 22:43:39 -0700 Subject: [PATCH 2/4] [NVBUG-6618098][test] Unwaive fixed PP4 stats test Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 8f54d011ddaf..12fdf46312e9 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -171,7 +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: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) From ccb2a9cbade16b8184a11c83c3d63462d9a7bd31 Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Tue, 18 Aug 2026 03:36:31 -0700 Subject: [PATCH 3/4] [NVBUG-6618102][test] Unwaive fixed H100 gpu4 LLM API suite The broad gpu4 waiver was added in the same failure batch as the two exact PP test waivers fixed by this PR. Remove it so the complete H100 marker selection runs again. Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 12fdf46312e9..0ce8f7585e4c 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -170,7 +170,6 @@ full:DGX_B200/disaggregated/test_disaggregated.py::test_disaggregated_gpt_oss_12 full:DGX_B200/perf/test_perf_sanity.py::test_e2e[aggr_upload-gemma4_26b_a4b_nvfp4_blackwell-gemma4_26b_a4b_nvfp4_tp1_1k1k] SKIP (https://nvbugs/6571410) 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: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) From 6768fe11979a26e4fa4714e3cc9935f8c146254e Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Tue, 18 Aug 2026 10:14:51 -0700 Subject: [PATCH 4/4] Revert "[NVBUG-6618102][test] Unwaive fixed H100 gpu4 LLM API suite" This reverts commit ccb2a9cbade16b8184a11c83c3d63462d9a7bd31. Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- tests/integration/test_lists/waives.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index 0ce8f7585e4c..12fdf46312e9 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -170,6 +170,7 @@ full:DGX_B200/disaggregated/test_disaggregated.py::test_disaggregated_gpt_oss_12 full:DGX_B200/perf/test_perf_sanity.py::test_e2e[aggr_upload-gemma4_26b_a4b_nvfp4_blackwell-gemma4_26b_a4b_nvfp4_tp1_1k1k] SKIP (https://nvbugs/6571410) 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: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)