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

[trtllm-agent]

Summary

• Root cause: The stable-greedy fast path's compact 1-D new_tokens_host was marked only on the outer SampleStateTorch, but PP transfers only sample_state.host, so non-final ranks rebuilt the state with single_step_greedy=False and indexed the flat buffer as 3-D [step][slot][beam].
• Fix: Move the layout marker onto SampleStateTensorsHostTorch (the object actually transferred) so the buffer is self-describing and send/recv stay symmetric with no wire-format change, and dispatch on either copy; plus an early return in the greedy branch of sampler_strategy.sample() that clears the blocking pre-existing mypy error under both the CI pin and the ambient tool.
• 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/6618100

Dev Engineer Review

• Moved single_step_greedy from SampleStateTorch to SampleStateTensorsHostTorch.
• Updated state access and storage to use the transferred host state.
• Preserved wire-format compatibility.
• Allowed dispatch to recognize the marker in either state location.
• Added an early return for greedy sampling to resolve the mypy error.
• The changes are small and consistent with the existing API usage.
• No configuration or test-list files changed.

QA Engineer Review

• Modified sampler tests to read single_step_greedy from SampleStateTensorsHostTorch.
• The provided summary does not identify the modified test function.
• Coverage in tests/integration/test_lists/ is not provided.
• Related test coverage includes test_tinyllama_logits_processor_tp2pp2.
• 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: 32faad84-cd17-4904-8bf4-dab1655d9c49

📥 Commits

Reviewing files that changed from the base of the PR and between 8efd46e and 0d6d0e9.

📒 Files selected for processing (3)

• tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
• tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.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 sampler moves single_step_greedy into host tensor state, updates state propagation and request handling, adjusts greedy return behavior, and updates related tests.

Changes

Greedy state alignment

Layer / File(s)	Summary
Host state contract and propagation tensorrt_llm/_torch/pyexecutor/sampler/sampler.py, tests/unittest/_torch/sampler/test_torch_sampler.py	single_step_greedy now belongs to SampleStateTensorsHostTorch. Sampling and request updates use the host state. Tests use the updated location.
Greedy return path tensorrt_llm/_torch/pyexecutor/sampler/sampler_strategy.py	The greedy branch directly returns tokens, optional probabilities, and None for temperature.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to 0d6d0

The PR makes a localized state-layout fix and does not introduce an actionable merge-blocking risk; it is merge-ready after normal checks and review.

Suggested reviewers: bowenfu, tongyuantongyu

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the bug fix and the main change to move the layout marker to SampleStateTensorsHostTorch.
Description check	✅ Passed	The description explains the root cause, fix, affected behavior, tests, and bug link, but does not include the full checklist template.
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.}

[HandongLi-01]

Dup of #17828, close this PR.