[https://nvbugs/6618106][fix] Carry single_step_greedy on SampleStateTensorsHostTorch by trtllm-agent · Pull Request #17828 · NVIDIA/TensorRT-LLM

trtllm-agent · 2026-08-17T11:58:33Z

Summary

Root cause: The stable-greedy fast path's compact 1-D new_tokens_host layout was flagged only on the outer SampleStateTorch, which PP's _ring_broadcast_sample_state never transmits, so non-final ranks paired the flat buffer with the 3-D reader.
Fix: Carry single_step_greedy on SampleStateTensorsHostTorch (the object actually pickled) and dispatch on either copy, keeping send/recv symmetric with no wire-format change; clear the pre-existing sampler_strategy mypy error via an early return that is clean under both the ambient and CI-pinned mypy; drop this bug's waives.txt line.
Original test: pytest tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_tinyllama_logits_processor_tp2pp2 -v
Automated fix generated by repair-bot

Test plan

Verify fix on the same GPU type as the original failure
Check for regressions in related tests

Links

Bug: https://nvbugs/6618106

Dev Engineer Review

single_step_greedy now travels with SampleStateTensorsHostTorch.
Non-final pipeline-parallel ranks can interpret compact 1-D new_tokens_host buffers correctly.
Send and receive behavior remains symmetric.
The wire format remains unchanged.
The early return resolves the existing sampler_strategy.py mypy error.
The related H100 waiver entries were removed.
No performance, error-handling, API consistency, or test-list format issues are evident.

QA Engineer Review

Modified tests/unittest/_torch/sampler/test_torch_sampler.py.
The sampler test setup and result handling now use SampleStateTensorsHostTorch.single_step_greedy.
The related test is tests/unittest/llmapi/test_llm_multi_gpu_pytorch.py::test_tinyllama_logits_processor_tp2pp2.
The related waiver entry was removed from tests/integration/test_lists/waives.txt.
The first downstream pipeline failed.
The second CI run completed successfully, including the downstream L0_MergeRequest_PR pipeline.

Verdict: sufficient

…ransmitted 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>

coderabbitai · 2026-08-17T12:02:18Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: dd1ffc47-8cb8-4f4b-8e20-3b13be4eb1a7

📥 Commits

Reviewing files that changed from the base of the PR and between ccb2a9c and 6768fe1.

📒 Files selected for processing (1)

tests/integration/test_lists/waives.txt

🚧 Files skipped from review as they are similar to previous changes (1)

tests/integration/test_lists/waives.txt

Included review availability: Your plan includes up to 12 reviews per rolling hour; 11 remain after this review.

Walkthrough

The sampler stores single_step_greedy in SampleStateTensorsHostTorch, reads it during request updates, returns directly from the greedy path, and updates related tests and integration waivers.

Changes

Sampler state and greedy-path updates

Layer / File(s)	Summary
Host-state propagation `tensorrt_llm/_torch/pyexecutor/sampler/sampler.py`	The `single_step_greedy` field moves from `SampleStateTorch` to `SampleStateTensorsHostTorch`. `sample_async` stores the field in host state, and `update_requests` reads it there.
Sampling return and validation `tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py`, `tests/unittest/_torch/sampler/test_torch_sampler.py`, `tests/integration/test_lists/waives.txt`	The greedy path returns `(tokens, softmax, None)` directly. Unit tests use the host-side field, and two DGX H100 waiver records are removed.

Estimated code review effort: 1 (Trivial) | ~5 minutes

Merge Risk: ⚪ Minimal · up to 6768f

The change carries the greedy-layout flag with the state that is transmitted between ranks, keeping buffer interpretation symmetric; no actionable merge-blocking risk remains beyond normal checks and review.

Suggested reviewers: zhaoyuanh-nvidia, yihuilu512

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 33.33% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title uses the required NVBugs and fix format and clearly describes the primary change.
Description check	✅ Passed	The description explains the root cause, fix, bug reference, and relevant test plan, although it uses different section headings than the template.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands.}

cascade812 · 2026-08-17T18:41:11Z

/bot run

tensorrt-cicd · 2026-08-17T18:47:22Z

PR_Github #66835 [ run ] triggered by Bot. Commit: f49d5b7 Link to invocation

tensorrt-cicd · 2026-08-17T21:38:53Z

PR_Github #66835 [ run ] completed with state SUCCESS. Commit: f49d5b7
/LLM/main/L0_MergeRequest_PR pipeline #54395 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

cascade812 · 2026-08-18T05:26:05Z

/bot run

tensorrt-cicd · 2026-08-18T05:33:15Z

PR_Github #66988 [ run ] triggered by Bot. Commit: f49d5b7 Link to invocation

Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com>

tensorrt-cicd · 2026-08-18T10:00:49Z

PR_Github #66988 [ run ] completed with state SUCCESS. Commit: f49d5b7
/LLM/main/L0_MergeRequest_PR pipeline #54532 completed with status: 'SUCCESS'

CI Report

Link to invocation

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>

This reverts commit ccb2a9c. Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com>

cascade812 · 2026-08-18T17:21:02Z

/bot skip --comment "revert last unwaive, previous CI success"

tensorrt-cicd · 2026-08-18T17:27:13Z

PR_Github #67123 [ skip ] triggered by Bot. Commit: 6768fe1 Link to invocation

tensorrt-cicd · 2026-08-18T17:37:01Z

PR_Github #67123 [ skip ] completed with state SUCCESS. Commit: 6768fe1
Skipping testing for commit 6768fe1

Link to invocation

trtllm-agent requested a review from a team as a code owner August 17, 2026 11:58

trtllm-agent requested review from YihuiLu512 and zhaoyuanh-nvidia August 17, 2026 11:58

trtllm-agent assigned cascade812 Aug 17, 2026

github-actions Bot assigned trtllm-agent Aug 17, 2026

cascade812 approved these changes Aug 17, 2026

View reviewed changes

cascade812 changed the title ~~[https://nvbugs/6618106][fix] Carry single_step_greedy on SampleStateTensorsHostTorch (the object actually…~~ [https://nvbugs/6618106][fix] Carry single_step_greedy on SampleStateTensorsHostTorch Aug 17, 2026

[NVBUG-6618098][test] Unwaive fixed PP4 stats test

ed860b2

Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com>

coderabbitai Bot mentioned this pull request Aug 18, 2026

[https://nvbugs/6627041][fix] Move single_step_greedy onto SampleStateTensorsHostTorch (the pickled… #17875

Closed

2 tasks

cascade812 added 2 commits August 18, 2026 03:36

Revert "[NVBUG-6618102][test] Unwaive fixed H100 gpu4 LLM API suite"

6768fe1

This reverts commit ccb2a9c. Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com>

cascade812 enabled auto-merge (squash) August 18, 2026 17:27

cascade812 merged commit b36faf1 into NVIDIA:main Aug 18, 2026
8 checks passed

This was referenced Aug 18, 2026

[https://nvbugs/6428124][fix] Include use_host_stop_criteria in the PP send/recv payload in… #16134

Open

[https://nvbugs/6619882][test] Unwaive unittest/_torch/sampler on DGX_B200 #17946

Merged

HandongLi-01 mentioned this pull request Aug 19, 2026

[https://nvbugs/6618100][fix] Move the layout marker onto SampleStateTensorsHostTorch (the object actually… #17826

Closed

2 tasks

cascade812 mentioned this pull request Aug 21, 2026

[https://nvbugs/6618102][chore] Unwaive multi gpu pp tests #17989

Merged

1 task

Navigation Menu

Uh oh!

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

[https://nvbugs/6618106][fix] Carry single_step_greedy on SampleStateTensorsHostTorch#17828
cascade812 merged 4 commits into
NVIDIA:mainfrom
tensorrt-cicd:repair-bot-bug6618106

trtllm-agent commented Aug 17, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

coderabbitai Bot commented Aug 17, 2026 •

edited

Loading

❌ Failed checks (1 warning)

Uh oh!

cascade812 commented Aug 17, 2026

Uh oh!

tensorrt-cicd commented Aug 17, 2026

Uh oh!

tensorrt-cicd commented Aug 17, 2026

Uh oh!

cascade812 commented Aug 18, 2026

Uh oh!

tensorrt-cicd commented Aug 18, 2026

Uh oh!

tensorrt-cicd commented Aug 18, 2026

Uh oh!

cascade812 commented Aug 18, 2026

Uh oh!

tensorrt-cicd commented Aug 18, 2026

Uh oh!

tensorrt-cicd commented Aug 18, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

trtllm-agent commented Aug 17, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test plan

Links

Dev Engineer Review

QA Engineer Review

Uh oh!

coderabbitai Bot commented Aug 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

cascade812 commented Aug 17, 2026

Uh oh!

tensorrt-cicd commented Aug 17, 2026

Uh oh!

tensorrt-cicd commented Aug 17, 2026

Uh oh!

cascade812 commented Aug 18, 2026

Uh oh!

tensorrt-cicd commented Aug 18, 2026

Uh oh!

tensorrt-cicd commented Aug 18, 2026

Uh oh!

cascade812 commented Aug 18, 2026

Uh oh!

tensorrt-cicd commented Aug 18, 2026

Uh oh!

tensorrt-cicd commented Aug 18, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

trtllm-agent commented Aug 17, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented Aug 17, 2026 •

edited

Loading