[https://nvbugs/6428069][fix] Move use_host_stop_criteria from SampleStateTorch onto…

[trtllm-agent]

Summary

• Root cause: use_host_stop_criteria lived on outer SampleStateTorch, not on the host object propagated across PP ring send/recv, so non-last PP ranks kept flag=False and dereferenced a None finish_reasons.
• Fix: Move use_host_stop_criteria from SampleStateTorch onto SampleStateTensorsHostTorch; update the one reader and the unit test; remove the eight nvbugs/6428069 waivers.
• 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/6428069

Summary by CodeRabbit

• Bug Fixes

  • Improved how stop conditions are tracked during generation, making request handling more consistent and reliable.
  • Adjusted sample-state handling so host-side criteria are applied in the right place.

• Tests

  • Updated test coverage to match the new stop-criteria behavior.
  • Refined integration test waivers to reflect the current expected set of scenarios.

[coderabbitai]

📝 Walkthrough

Walkthrough

The use_host_stop_criteria boolean field is relocated from SampleStateTorch to SampleStateTensorsHostTorch, with consumer code (update_requests(), sample_async()) and a unit test updated accordingly. Separately, waiver entries in waives.txt are narrowed/removed for several disaggregated serving test configurations.

Changes

Host stop criteria field relocation

Layer / File(s)	Summary
Move field and update consumers tensorrt_llm/_torch/pyexecutor/sampler/sampler.py, tests/unittest/_torch/sampler/test_torch_sampler.py	use_host_stop_criteria moves from SampleStateTorch to SampleStateTensorsHostTorch; update_requests() and sample_async() are updated to read/set it from state.host, and the corresponding unit test assertion order is updated to match.

Test waiver list updates

Layer / File(s)	Summary
Update disaggregated serving waivers tests/integration/test_lists/waives.txt	SKIP entries for TestLlama3_1_8BInstruct are narrowed to specific gen_tp/ctx_pp=2 combinations, test_ngram, and tp1pp2 cases, and two TinyLlama disaggregated waiver entries are removed.

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

Possibly related PRs

• NVIDIA/TensorRT-LLM#16103: Both PRs modify overlapping SKIP/waiver entries in tests/integration/test_lists/waives.txt for disaggregated serving test configurations.

Suggested reviewers: yuxianq, dongxuy04

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is specific, concise, and matches the main change of moving use_host_stop_criteria to host state.
Description check	✅ Passed	The description clearly covers the root cause, fix, tests, and bug link, even though it does not use the template headings exactly.
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

⚔️ Resolve merge conflicts

• Resolve merge conflict in branch repair-bot-bug6428069

---

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

[coderabbitai]

🧹 Nitpick comments (1)

tests/unittest/_torch/sampler/test_torch_sampler.py (1)

932-934: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

LGTM for the assertion update itself, matching the new field location.

One coverage gap worth flagging: this test exercises sample_async/update_requests within a single sampler instance and doesn't cover the actual PP send/recv boundary that caused bug 6428069 (non-last ranks receiving state.host without the flag). Consider a lightweight unit test that pickles/round-trips (or otherwise simulates the PP transfer of) SampleStateTensorsHostTorch to assert use_host_stop_criteria survives, complementing the disaggregated PP integration tests already re-enabled by this PR.

As per path instructions for tests/**: "suggest concrete list file names and whether coverage is sufficient, insufficient, or needs follow-up outside the PR."

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unittest/_torch/sampler/test_torch_sampler.py` around lines 932 - 934,
Coverage is insufficient for the PP transfer path: the current test only checks
`sample_async` and `update_requests` inside one `TorchSampler` instance, so it
does not verify that `SampleStateTensorsHostTorch.use_host_stop_criteria`
survives the send/recv boundary. Add a lightweight unit test in
`tests/unittest/_torch/sampler/test_torch_sampler.py` that round-trips or
simulates serialization for `SampleStateTensorsHostTorch` and asserts
`use_host_stop_criteria` remains set on the received `state.host`. This should
complement, not replace, the existing disaggregated PP integration tests already
covered elsewhere in the PR.

Source: Path instructions

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@tests/unittest/_torch/sampler/test_torch_sampler.py`:
- Around line 932-934: Coverage is insufficient for the PP transfer path: the
current test only checks `sample_async` and `update_requests` inside one
`TorchSampler` instance, so it does not verify that
`SampleStateTensorsHostTorch.use_host_stop_criteria` survives the send/recv
boundary. Add a lightweight unit test in
`tests/unittest/_torch/sampler/test_torch_sampler.py` that round-trips or
simulates serialization for `SampleStateTensorsHostTorch` and asserts
`use_host_stop_criteria` remains set on the received `state.host`. This should
complement, not replace, the existing disaggregated PP integration tests already
covered elsewhere in the PR.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: a8365454-2569-4bac-8dab-fba3226a42dd

📥 Commits

Reviewing files that changed from the base of the PR and between 5ce293f and 0b942e8.

📒 Files selected for processing (3)

• tensorrt_llm/_torch/pyexecutor/sampler/sampler.py
• tests/integration/test_lists/waives.txt
• tests/unittest/_torch/sampler/test_torch_sampler.py

💤 Files with no reviewable changes (1)

• tests/integration/test_lists/waives.txt