[None][perf] Move greedy stop checks to host by mingyangHao · Pull Request #15920 · NVIDIA/TensorRT-LLM

mingyangHao · 2026-07-03T10:35:53Z

Summary by CodeRabbit

New Features
- Added a host-side stop-criteria path for fast greedy sampling in eligible single-token, single-beam requests.
- Improved request handling so sampled tokens can be finalized directly on the host in supported cases.
Bug Fixes
- Reduced unnecessary device-side finish-reason processing when host stop criteria can be used.
- Adjusted token placement logic to work correctly with the new fast path.
Tests
- Added coverage for host stop-criteria behavior and fallback cases where it should remain disabled.

Description

A focused benchmark on NVIDIA B300 with 64 requests, using the actual TorchSampler backend, showed:

Median sampler latency: 454.28 µs → 215.94 µs
Latency reduction: 238.35 µs (52.47%)
Sampler speedup: 2.10×

Test Coverage

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

mingyangHao · 2026-07-03T10:36:20Z

/bot run --disable-fail-fast

coderabbitai · 2026-07-03T10:40:52Z

📝 Walkthrough

Walkthrough

This PR adds an optional host-side stop-criteria evaluation path for single-token, single-beam fast-greedy sampling in TorchSampler. A new use_host_stop_criteria flag is computed, propagated through SampleStateTorch, _process_requests, sample_async, and update_requests, and covered by new unit tests.

Changes

Host stop criteria for fast-greedy sampling

Layer / File(s)	Summary
State contract and eligibility check `tensorrt_llm/_torch/pyexecutor/sampler.py`	Adds `use_host_stop_criteria` field to `SampleStateTorch` and a new `_can_use_host_stop_criteria` helper that checks fast-greedy eligibility, single-token/single-beam constraints, and absence of stop words/draft requests.
_process_requests computation and propagation `tensorrt_llm/_torch/pyexecutor/sampler.py`	Computes `use_fast_greedy_path`/`use_host_stop_criteria`, skips building `seq_lens_host` when enabled, derives scatter destination indices from `seq_slots_cuda`, adds assertions for the non-fast-greedy path, and extends the return tuple at all return points.
sample_async and update_requests wiring `tensorrt_llm/_torch/pyexecutor/sampler.py`	`sample_async` captures and forwards the new flag, skips device finish-reasons setup when enabled; `update_requests` branches to add sampled tokens and call `_handle_stop_criteria` on host instead of device-based draft-token handling.
Unit tests `tests/unittest/_torch/sampler/test_torch_sampler.py`	Adds parametrized tests for host stop-criteria fast path (state flags, skipped device finish-reasons, expected `finish_by` calls) and fallback cases confirming `_can_use_host_stop_criteria` returns `False` for stop-words/draft requests.

Estimated code review effort: 4 (Complex) | ~45 minutes

Sequence Diagram(s)

sequenceDiagram
participant SampleAsync as sample_async
participant ProcessRequests as _process_requests
participant SampleState as SampleStateTorch
participant UpdateRequests as update_requests
participant StopCriteria as _handle_stop_criteria

SampleAsync->>ProcessRequests: invoke sampling
ProcessRequests->>ProcessRequests: compute use_host_stop_criteria
ProcessRequests-->>SampleAsync: return tokens, use_host_stop_criteria
SampleAsync->>SampleState: set use_host_stop_criteria flag
SampleAsync->>UpdateRequests: pass sample state
UpdateRequests->>UpdateRequests: check state.use_host_stop_criteria
alt host stop criteria enabled
UpdateRequests->>UpdateRequests: add sampled token to request
UpdateRequests->>StopCriteria: evaluate stop criteria on host
else disabled
UpdateRequests->>UpdateRequests: run device finish-reasons + draft-token processing
end

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The description includes a benchmark result, but the Test Coverage section is empty and the what-and-why details are incomplete.	Add a brief what-and-why summary and list the tests covering the new host stop-criteria path.

✅ Passed checks (4 passed)

Check name	Status	Explanation
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.
Title check	✅ Passed	The title follows the required [None][perf] format and clearly summarizes moving greedy stop checks to the host.

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

tensorrt-cicd · 2026-07-03T10:42:37Z

PR_Github #57453 [ run ] triggered by Bot. Commit: cdbd8a3 Link to invocation

tensorrt-cicd · 2026-07-03T15:51:25Z

PR_Github #57453 [ run ] completed with state SUCCESS. Commit: cdbd8a3
/LLM/main/L0_MergeRequest_PR pipeline #46190 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

Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>

mingyangHao · 2026-07-06T01:49:10Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-06T01:54:34Z

PR_Github #57647 [ run ] triggered by Bot. Commit: 60e8d0e Link to invocation

tensorrt-cicd · 2026-07-06T07:31:20Z

PR_Github #57647 [ run ] completed with state SUCCESS. Commit: 60e8d0e
/LLM/main/L0_MergeRequest_PR pipeline #46369 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com> Co-authored-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>

Remove the PP sampler waivers introduced by NVIDIA#16103, NVIDIA#16105, NVIDIA#16127, and NVIDIA#16146. Also remove overlapping PP and gpu2 waivers from NVIDIA#16109 so the restored pre-NVIDIA#15920 sampler path is exercised in CI. Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>

This reverts commit 4f1c3cd. Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>

Remove the PP sampler waivers introduced by NVIDIA#16103, NVIDIA#16105, NVIDIA#16127, and NVIDIA#16146. Also remove overlapping PP and gpu2 waivers from NVIDIA#16109, plus the PP2 waiver added by NVIDIA#16169, so the restored pre-NVIDIA#15920 sampler path is exercised in CI. Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>

…16163) Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com> Co-authored-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>

Reland the host-side greedy stop-criteria fast path from NVIDIA#15920 after its temporary revert in NVIDIA#16163. Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>

Remove the PP sampler waivers introduced by NVIDIA#16103, NVIDIA#16105, NVIDIA#16127, and NVIDIA#16146. Also remove overlapping PP and gpu2 waivers from NVIDIA#16109, plus the PP2 waiver added by NVIDIA#16169, so the restored pre-NVIDIA#15920 sampler path is exercised in CI. Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>

PR NVIDIA#15920 introduced a fast path in TorchSampler that skips writing finish_reasons to host when all requests are greedy, single-beam, and have no stop words. In pipeline-parallel mode with disable_overlap_scheduler=False, only the last PP rank runs sample_async and knows the fast path was taken; non-last PP ranks construct SampleState in _forward_step_inter_pp with the default use_host_stop_criteria=False. The ring-broadcast payload sent only sample_state.host, so non-last ranks ended up with finish_reasons=None but use_host_stop_criteria=False, falling through to process_draft_tokens -> finish_if_reason which indexed an empty list and raised IndexError. Include use_host_stop_criteria in the PP ring-broadcast payload so all ranks agree on which update_requests branch to take. Guard with getattr/hasattr so SampleStateTRTLLM (which lacks this field) is unaffected. Remove the waiver for the affected test. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>

The TorchSampler IndexError described in nvbug 6435097 was introduced by PR NVIDIA#15920 ("Move greedy stop checks to host") and cleanly reverted by PR NVIDIA#16163 (commit 9a8ec05). PR NVIDIA#16163 is already present at HEAD, so the underlying code bug is fixed. The waiver at nvbugs/6427411 for test_ctx_pp_gen_tp_asymmetric[GSM8K-gen_tp=1-ctx_pp=2] was added by PRs NVIDIA#16103 / NVIDIA#16105 / NVIDIA#16127 before the revert landed and was not cleaned up afterward. Verified locally on this HEAD (e523b43): pytest tests/integration/defs/accuracy/test_disaggregated_serving.py::\ TestLlama3_1_8BInstruct::test_ctx_pp_gen_tp_asymmetric[GSM8K-gen_tp=1-ctx_pp=2] -> PASSED, GSM8K accuracy = 73.465 (threshold 70.997). Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>

mingyangHao requested a review from a team as a code owner July 3, 2026 10:35

mingyangHao requested a review from achartier July 3, 2026 10:35

github-actions Bot assigned mingyangHao Jul 3, 2026

achartier approved these changes Jul 5, 2026

View reviewed changes

[None][perf] Move greedy stop checks to host

60e8d0e

Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>

mingyangHao force-pushed the user/mingyangh/torch-sampler-host-stop branch from cdbd8a3 to 60e8d0e Compare July 6, 2026 01:47

longlee0622 approved these changes Jul 7, 2026

View reviewed changes

longlee0622 merged commit 4f1c3cd into NVIDIA:main Jul 7, 2026
12 of 13 checks passed

mingyangHao mentioned this pull request Jul 8, 2026

[None][fix] Propagate host stop criteria across PP ranks #16096

Open

1 task

This was referenced Jul 8, 2026

[https://nvbugs/6427411][chore] Waive remaining PP sampler regressions #16127

Merged

[https://nvbugs/6428057][fix] Propagate host stop criteria across PP ranks #16131

Closed

trtllm-agent mentioned this pull request Jul 8, 2026

[https://nvbugs/6428063][fix] Include use_host_stop_criteria in the PP ring-broadcast payload (sender uses… #16136

Open

2 tasks

chienchunhung mentioned this pull request Jul 8, 2026

[https://nvbugs/6427411][chore] Waive additional PP sampler regressions #16146

Merged

mingyangHao added a commit to mingyangHao/TensorRT-LLM that referenced this pull request Jul 9, 2026

Revert "[None][perf] Move greedy stop checks to host (NVIDIA#15920)"

6fcfd81

This reverts commit 4f1c3cd. Signed-off-by: Mingyang Hao <200044211+mingyangHao@users.noreply.github.com>

trtllm-agent mentioned this pull request Jul 11, 2026

[https://nvbugs/6435097][fix] Remove the single stale nvbugs/6427411 waiver line for… #16266

Open

2 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[None][perf] Move greedy stop checks to host - #15920

[None][perf] Move greedy stop checks to host#15920
longlee0622 merged 1 commit into
NVIDIA:mainfrom
mingyangHao:user/mingyangh/torch-sampler-host-stop

mingyangHao commented Jul 3, 2026 •

edited

Loading

Uh oh!

mingyangHao commented Jul 3, 2026

Uh oh!

coderabbitai Bot commented Jul 3, 2026 •

edited

Loading

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (1 warning)

Uh oh!

tensorrt-cicd commented Jul 3, 2026

Uh oh!

tensorrt-cicd commented Jul 3, 2026

Uh oh!

mingyangHao commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

mingyangHao commented Jul 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

mingyangHao commented Jul 3, 2026

Uh oh!

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

Uh oh!

Walkthrough

Changes

Sequence Diagram(s)

❌ Failed checks (1 warning)

Uh oh!

tensorrt-cicd commented Jul 3, 2026

Uh oh!

tensorrt-cicd commented Jul 3, 2026

Uh oh!

mingyangHao commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

tensorrt-cicd commented Jul 6, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

mingyangHao commented Jul 3, 2026 •

edited

Loading

coderabbitai Bot commented Jul 3, 2026 •

edited

Loading