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

[trtllm-agent]

Summary

• Root cause: single_step_greedy lived on SampleStateTorch, which the PP ring never transports, so non-final ranks defaulted it to False and read the compact 1-D new_tokens_host as 3-D.
• Fix: Move single_step_greedy onto SampleStateTensorsHostTorch (the pickled object) and read it via state.host, giving the layout flag a single owner that travels with the tensors.
• Original test: pytest tests/integration/defs/perf/test_perf.py --perf --test-list=.repair-bot/perf_test_list.txt --output-dir=build/perf_output -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/6627041

Dev Engineer Review

• single_step_greedy moved to SampleStateTensorsHostTorch.
• update_requests now reads the flag from state.host.
• This keeps the flag with host new_tokens data and prevents incorrect interpretation of compact 1-D data on non-final ranks.
• The change is limited to the sampler state model and preserves the default value.
• No configuration or test-list files changed.

QA Engineer Review

• Modified test code: tests/unittest/_torch/sampler/test_torch_sampler.py.
• The test helper now reads single_step_greedy from sample_state.host.
• The existing sampler regression coverage was updated to match the new state ownership.
• The test is not listed in the provided test-list context.
• Verdict: needs follow-up.

[coderabbitai]

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: a3640b33-7be9-41f9-8d59-5db8b4c35d90

📥 Commits

Reviewing files that changed from the base of the PR and between 0eeda34 and d2e8419.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
• tests/unittest/_torch/sampler/test_torch_sampler.py

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

---

Walkthrough

The change moves single_step_greedy from SampleStateTorch to SampleStateTensorsHostTorch. Sampling, request updates, and tests now access the flag through the host state.

Changes

Sampler state update

Layer / File(s)	Summary
State ownership and propagation tensorrt_llm/_torch/pyexecutor/sampler/sampler.py	SampleStateTensorsHostTorch now owns single_step_greedy. sample_async stores the flag in the host state, and update_requests reads it from there.
Sampler test updates tests/unittest/_torch/sampler/test_torch_sampler.py	Tests initialize and inspect single_step_greedy through sample_state.host.

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

Merge Risk: ⚪ Minimal · up to d2e84

This localized fix is merge-ready after normal checks; no actionable merge-blocking risk remains.

Possibly related PRs

• NVIDIA/TensorRT-LLM#17826: Implements the same sampler state field relocation and test updates.
• NVIDIA/TensorRT-LLM#17828: Updates the same field usage in sampler.py and its tests.
• NVIDIA/TensorRT-LLM#17818: Relocates single_step_greedy and updates its propagation and consumption.

Suggested reviewers: bowenfu

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the bug fix and the field move described by the pull request.
Description check	✅ Passed	The description explains the root cause and fix, identifies tests, and links the bug; the repository checklist is not included.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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]

Duplicate PR with #17828