[https://nvbugs/6427240][fix] Reserve MTP draft tokens in scheduler for one-model speculative decoding - #16101
Conversation
jhaotingc
changed the title
jhaotingc
changed the title
jhaotingc
changed the title
jhaotingc
changed the title
📝 WalkthroughWalkthroughThis PR adds a conditional branch in ChangesDraft Token Initialization
Estimated code review effort: 2 (Simple) | ~10 minutes 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/py_executor.py (1)
2980-3009: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueConsider extracting the shared draft-token reservation loop.
The state-filter +
draft_tokensassignment loop (lines 2981-2988 and 3004-3009) is duplicated between the two-model and one-model branches. Extracting a small helper (e.g._reserve_draft_token_budget(requests, count)) would reduce duplication and keep the two normalization paths in sync going forward.♻️ Example helper extraction
+ def _reserve_draft_token_budget(self, num_draft_tokens: int) -> None: + for request in self.active_requests: + if request.state not in ( + LlmRequestState.GENERATION_IN_PROGRESS, + LlmRequestState.DISAGG_GENERATION_INIT): + continue + request.draft_tokens = [0] * num_draft_tokens + if self.drafter is not None: ... - for request in self.active_requests: - if request.state not in ( - LlmRequestState.GENERATION_IN_PROGRESS, - LlmRequestState.DISAGG_GENERATION_INIT): - continue - request.draft_tokens = [ - 0 - ] * self.max_total_draft_tokens if self.max_total_draft_tokens > 0 else [] + self._reserve_draft_token_budget( + self.max_total_draft_tokens if self.max_total_draft_tokens > 0 else 0) ... elif self.model_engine.is_spec_decode and self.max_total_draft_tokens > 0: ... - for request in self.active_requests: - if request.state not in ( - LlmRequestState.GENERATION_IN_PROGRESS, - LlmRequestState.DISAGG_GENERATION_INIT): - continue - request.draft_tokens = [0] * self.max_total_draft_tokens + self._reserve_draft_token_budget(self.max_total_draft_tokens)🤖 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 `@tensorrt_llm/_torch/pyexecutor/py_executor.py` around lines 2980 - 3009, The state-filtered draft-token reservation loop is duplicated in both the two-model and one-model speculative decoding branches. Extract the shared logic in py_executor.py into a small helper (for example, a method near _prepare_draft_requests or the surrounding scheduling code) that iterates over self.active_requests, filters by LlmRequestState.GENERATION_IN_PROGRESS and LlmRequestState.DISAGG_GENERATION_INIT, and assigns request.draft_tokens based on the requested count; then call that helper from both branches so the reservation behavior stays in sync.
🤖 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 `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 2980-3009: The state-filtered draft-token reservation loop is
duplicated in both the two-model and one-model speculative decoding branches.
Extract the shared logic in py_executor.py into a small helper (for example, a
method near _prepare_draft_requests or the surrounding scheduling code) that
iterates over self.active_requests, filters by
LlmRequestState.GENERATION_IN_PROGRESS and
LlmRequestState.DISAGG_GENERATION_INIT, and assigns request.draft_tokens based
on the requested count; then call that helper from both branches so the
reservation behavior stays in sync.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 3cedd69c-84b2-4bf3-9fd2-e73ab4fb1273
📒 Files selected for processing (1)
tensorrt_llm/_torch/pyexecutor/py_executor.py
jhaotingc
changed the title
jhaotingc
changed the title
3685480 to
6f3bb4a
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #58152 [ run ] triggered by Bot. Commit: |
|
/bot run |
|
PR_Github #58265 [ run ] triggered by Bot. Commit: |
|
PR_Github #58152 [ run ] completed with state
|
|
PR_Github #58265 [ run ] completed with state
|
6f3bb4a to
2e6e8bd
Compare
|
/bot run |
|
PR_Github #58327 [ run ] triggered by Bot. Commit: |
2e6e8bd to
9e45af3
Compare
|
PR_Github #58327 [ run ] completed with state
|
9e45af3 to
a6eb3b5
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #58481 [ run ] triggered by Bot. Commit: |
|
CI all passed, blocked by #16208 |
|
/bot kill |
|
PR_Github #58541 [ kill ] triggered by Bot. Commit: |
|
PR_Github #58481 [ run ] completed with state
|
|
PR_Github #58541 [ kill ] completed with state |
a6eb3b5 to
057a3bb
Compare
|
/bot run |
057a3bb to
2e11394
Compare
…or one-model speculative decoding
One-model speculative decoding (MTP / Eagle3 one-model) has no separate
drafter, so get_spec_drafter() returns None and the `if self.drafter is not
None` block in _prepare_and_schedule_batch is skipped -- including the loop
that sets `request.draft_tokens = [0] * max_total_draft_tokens` for generation
requests. That loop is what makes the C++ micro-batch scheduler budget each
generation request as beam_width + getNumDraftTokens(). Without it, under
enable_chunked_prefill + the overlap scheduler, the scheduler under-reserves
and the forward's uniform (1 + runtime_draft_len) build overshoots
max_num_tokens, aborting requests with an AssertionError (total_num_tokens >
max_num_tokens) in _prepare_tp_inputs.
Add an elif branch mirroring the two-model draft-token normalization for the
one-model path so scheduling reserves the correct budget. Using
self.max_total_draft_tokens keeps it consistent with dynamic draft_len_schedule.
model_engine is guarded (getattr ... is not None) first so partially-
constructed executors in unit tests do not raise AttributeError. The state
filter {GENERATION_IN_PROGRESS, DISAGG_GENERATION_INIT} matches the two-model
normalization, so this single fix covers both aggregated and disaggregated
(decode-worker) serving.
Add a regression test (test_py_executor.py) that drives
_prepare_and_schedule_batch on a one-model-MTP executor with real LlmRequests
in GENERATION_IN_PROGRESS (aggregated) and DISAGG_GENERATION_INIT (disagg)
states, asserting both are normalized to max_total_draft_tokens while
CONTEXT_INIT is left untouched. The test fails without the fix and passes with
it.
Verified on Qwen3.6-35B-A3B-NVFP4 (MTP num_nextn_predict_layers=3, chunked
prefill, overlap scheduler on): a concurrency sweep that previously produced
130 assertions at concurrency 32 now completes with zero assertions and zero
request errors.
Signed-off-by: Jhao-Ting Chen <jtchen0528@gmail.com>
2e11394 to
9cb5493
Compare
|
/bot run |
|
PR_Github #58548 [ run ] triggered by Bot. Commit: |
|
PR_Github #58550 [ run ] triggered by Bot. Commit: |
|
PR_Github #58548 [ run ] completed with state |
|
PR_Github #58550 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58662 [ run ] triggered by Bot. Commit: |
|
PR_Github #58662 [ run ] completed with state |
5 participants
Background / motivation
Serving a one-model MTP model (e.g. Qwen3.6-35B-A3B-NVFP4,
mtp_eagle_one_model) withenable_chunked_prefilland the overlap scheduler (disable_overlap_scheduler: false) aborts requests under load with:at
tensorrt_llm/_torch/pyexecutor/model_engine.py::_prepare_tp_inputs, surfaced to clients asRequestError. It reproduces at moderate concurrency with long (chunk-prefilled) prompts.Root cause
get_spec_drafter()returnsNoneforMTP_EAGLE_ONE_MODEL/ one-model speculative decoding, so theif self.drafter is not None:block in_prepare_and_schedule_batchis skipped — including the loop that setsrequest.draft_tokens = [0] * max_total_draft_tokensfor generation requests. That loop is what makes the C++ micro-batch scheduler (microBatchScheduler.cpp, costbeam_width + getNumDraftTokens()) reserve the draft-token budget. Without it the scheduler under-reserves, while the forward builds a uniform1 + runtime_draft_lenper gen request (overlap path); the chunked-prefill context chunk is then sized pastmax_num_tokens, tripping the assertion. Two-model speculative decoding runs the normalization and is unaffected.Summary
Add an
elifbranch that mirrors the two-model draft-token normalization for the one-model path (self.model_engine.is_spec_decode and self.max_total_draft_tokens > 0). It usesself.max_total_draft_tokens, so it stays consistent with dynamicdraft_len_schedule(a hardcoded max reservation would over-reserve). Scheduling now reservesbeam + max_total_draft_tokensper gen request, matching the forward.Impact
is_spec_decodeis False).Verification
Qwen3.6-35B-A3B-NVFP4, TP1 B200, MTP
num_nextn_predict_layers=3,enable_chunked_prefill, overlap scheduler on,max_num_tokens=8192,max_batch_size=512. aiperf concurrency sweep (NIAH long-prompt dataset, OSL 6144, ignore_eos):total_num_tokensassertsNotes
Opening as draft for approach review. Happy to add a regression unit test if desired.
Summary by CodeRabbit