[None][fix] Reserve one-engine speculative tokens before scheduling by yizhang-nv · Pull Request #15973 · NVIDIA/TensorRT-LLM

yizhang-nv · 2026-07-06T09:51:08Z

Description

One-engine speculative decoders keep their draft model inside ModelEngine, so PyExecutor.drafter is None. With the overlap scheduler, the next draft tokens can be produced by the previous batch before their count is copied to the C++ request. The scheduler can therefore budget a generation request as one token even though the model consumes the full 1 + K speculative step.

This can admit an oversized mixed context/generation batch. The observed MTP3 failure was total_num_tokens (8195) should be <= max_num_tokens (8192).

Move scheduler-facing draft-token preparation out of the executor-level Drafter branch and run it immediately before scheduling for every speculative decoder. The placeholder token values are not consumed by the model; their length reserves the correct micro-batch capacity.

Add a regression test that enters through _prepare_and_schedule_batch and verifies that a one-engine speculative request exposes all three draft tokens at the _schedule boundary.

Remove the NVBug 5989920 waiver for test_nvfp4_multi_gpus_piecewise_cuda_graph[mtp3_fp8kv_chunked] so CI runs the configuration that exposed this bug.

Test Coverage

pytest -s -q tests/unittest/_torch/executor/test_py_executor.py — 69 passed.
pytest -s -vv --tb=short tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_multi_gpus_piecewise_cuda_graph[mtp3_fp8kv_chunked] on 8x B200 — passed; MMLU 87.890, GSM8K 95.679.
pre-commit run --files tensorrt_llm/_torch/pyexecutor/py_executor.py tests/unittest/_torch/executor/test_py_executor.py tests/integration/test_lists/waives.txt — passed.

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.

Summary by CodeRabbit

Bug Fixes
- Improved handling of draft-token accounting during scheduling so generation requests are prepared more consistently when speculative decoding is involved.
- Added safer fallback behavior when speculative decoding components are unavailable, preventing unexpected scheduling state.
Tests
- Added coverage for mixed request scheduling scenarios to verify draft-token reservation and request processing behave as expected.

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>

Unwaive TestDeepSeekV32::test_nvfp4_multi_gpus_piecewise_cuda_graph[mtp3_fp8kv_chunked] (nvbugs/5989920) to re-enable it in CI. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>

coderabbitai · 2026-07-06T09:56:37Z

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: c06a88f8-fc1a-4ac0-83da-1003774777de

📥 Commits

Reviewing files that changed from the base of the PR and between 0044d5b and 9dda6c3.

📒 Files selected for processing (2)

tensorrt_llm/_torch/pyexecutor/py_executor.py
tests/unittest/_torch/executor/test_py_executor.py

📝 Walkthrough

Walkthrough

PyExecutor's _prepare_and_schedule_batch no longer inline-assigns request.draft_tokens; instead it calls a new _prepare_scheduler_draft_tokens() method that sets zero-filled placeholder draft tokens for active generation requests, skipping when no drafter or spec_config exists. A corresponding unit test was added.

Changes

Scheduler Draft Token Placeholder Refactor

Layer / File(s)	Summary
Extract draft token placeholder logic `tensorrt_llm/_torch/pyexecutor/py_executor.py`	Removes inline draft_tokens assignment and adds `_prepare_scheduler_draft_tokens()`, which sets zero-filled placeholder draft_tokens for active generation requests, skipping if neither a drafter nor spec_config exists.
Test coverage for draft token reservation `tests/unittest/_torch/executor/test_py_executor.py`	Imports `LlmRequest` and `SamplingConfig`, and adds a test verifying `_prepare_and_schedule_batch()` resets `draft_tokens` and sets `num_draft_tokens` correctly for generation requests.

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

Possibly related PRs

NVIDIA/TensorRT-LLM#12262: Related to _prepare_and_schedule_batch reserving scheduler-facing draft token placeholders, aligning with broader dynamic draft-length budgeting changes.

Suggested reviewers: syuoni, mikeiovine, brb-nv

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title matches the main change and follows the required [None][fix] summary format.
Description check	✅ Passed	The description includes the required Description, Test Coverage, and PR Checklist sections with concrete details.
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.}

yizhang-nv · 2026-07-07T03:13:18Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-07T03:19:15Z

PR_Github #57916 [ run ] triggered by Bot. Commit: de560e6 Link to invocation

achartier · 2026-07-07T04:21:35Z

+        if (self.drafter is None
+                and getattr(self.model_engine, 'spec_config', None) is None):
+            return
+


Should also early return if use_spec_decode is False

tensorrt-cicd · 2026-07-07T08:59:16Z

PR_Github #57916 [ run ] completed with state FAILURE. Commit: de560e6
/LLM/main/L0_MergeRequest_PR pipeline #46608 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: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>

yizhang-nv · 2026-07-08T05:31:45Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-08T05:37:24Z

PR_Github #58156 [ run ] triggered by Bot. Commit: e34fe34 Link to invocation

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>

yizhang-nv · 2026-07-08T05:38:06Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-08T05:44:13Z

PR_Github #58158 [ run ] triggered by Bot. Commit: b71f9f0 Link to invocation

tensorrt-cicd · 2026-07-08T05:47:44Z

PR_Github #58156 [ run ] completed with state ABORTED. Commit: e34fe34

Link to invocation

tensorrt-cicd · 2026-07-08T11:08:30Z

PR_Github #58158 [ run ] completed with state FAILURE. Commit: b71f9f0
/LLM/main/L0_MergeRequest_PR pipeline #46809 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

yizhang-nv · 2026-07-09T03:40:05Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-09T03:46:22Z

PR_Github #58362 [ run ] triggered by Bot. Commit: b71f9f0 Link to invocation

tensorrt-cicd · 2026-07-09T05:36:18Z

PR_Github #58362 [ run ] completed with state SUCCESS. Commit: b71f9f0
/LLM/main/L0_MergeRequest_PR pipeline #46988 completed with status: 'SUCCESS'

CI Report

Link to invocation

…upstream PR NVIDIA#15973) Squash of upstream PR NVIDIA#15973 (4 commits). One-engine speculative decoders (MTP / DSpark) keep the draft model inside ModelEngine, so PyExecutor.drafter is None; under the overlap scheduler the next draft tokens are produced by the previous batch and their count may not be copied to the C++ request before the next scheduling pass, so the scheduler can budget a gen request as 1 token though the model consumes 1+K -> oversized mixed context/generation batch (total_num_tokens > max_num_tokens). Run scheduler-facing draft-token reservation for every speculative decoder immediately before scheduling (_prepare_scheduler_draft_tokens); placeholder token values are not consumed, their length reserves the correct micro-batch capacity. Includes: skip when spec decode is disabled, handle lightweight executor fixtures, and unwaive the DeepSeekV3 nvfp4 mtp3_fp8kv_chunked test that exposed the bug. Signed-off-by: chungen28 <chung-en@deepinfra.com>

[None][fix] Reserve one-engine speculative tokens before scheduling

9dda6c3

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>

yizhang-nv requested a review from a team as a code owner July 6, 2026 09:51

yizhang-nv requested a review from achartier July 6, 2026 09:51

github-actions Bot assigned yizhang-nv Jul 6, 2026

[None][test] Unwaive DeepSeekV3 nvfp4 mtp3_fp8kv_chunked test

de560e6

Unwaive TestDeepSeekV32::test_nvfp4_multi_gpus_piecewise_cuda_graph[mtp3_fp8kv_chunked] (nvbugs/5989920) to re-enable it in CI. Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>

achartier approved these changes Jul 7, 2026

View reviewed changes

[None][fix] Skip scheduler draft tokens when spec decode is disabled

e34fe34

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>

[None][fix] Handle lightweight executor fixtures in draft token setup

b71f9f0

Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>

Uh oh!

Conversation

yizhang-nv commented Jul 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Test Coverage

PR Checklist

GitHub Bot Help

Summary by CodeRabbit

Uh oh!

coderabbitai Bot commented Jul 6, 2026

Walkthrough

Changes

Uh oh!

yizhang-nv commented Jul 7, 2026

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

achartier Jul 7, 2026

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

yizhang-nv commented Jul 8, 2026

Uh oh!

tensorrt-cicd commented Jul 8, 2026

Uh oh!

yizhang-nv commented Jul 8, 2026

Uh oh!

tensorrt-cicd commented Jul 8, 2026

Uh oh!

tensorrt-cicd commented Jul 8, 2026

Uh oh!

tensorrt-cicd commented Jul 8, 2026

Uh oh!

yizhang-nv commented Jul 9, 2026

Uh oh!

tensorrt-cicd commented Jul 9, 2026

Uh oh!

tensorrt-cicd commented Jul 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

yizhang-nv commented Jul 6, 2026 •

edited

Loading