[TRTLLM-13321][feat] Extend rejection sampling to one-model speculative decoding (MTP/DRAFT_TARGET/PARD/DFLASH)#15775
Merged
zhaoyangwang-nvidia merged 8 commits intoJul 16, 2026
Conversation
zhaoyangwang-nvidia
force-pushed
the
rejection-one-model-spec
branch
3 times, most recently
from
c8302fc to
3f93561
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #57463 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57463 [ run ] completed with state
|
zhaoyangwang-nvidia
force-pushed
the
rejection-one-model-spec
branch
2 times, most recently
from
ba30f91 to
fcd785a
Compare
zhaoyangwang-nvidia
requested review from
cascade812,
dongxuy04,
lfr-0531,
yechank-nvidia and
zhenhuaw-me
Contributor
📝 WalkthroughWalkthroughThis PR refactors one-model speculative decoding to add configurable LM head gather output, cache draft-to-target vocab mapping (d2t) on workers, unify draft token sampling and acceptance through new base methods, and introduce fail-closed rejection-sampling buffer validation. It also adds draft vocab size computation, updated config validation, developer documentation, and rejection-sampling tests across Eagle3, MTP, PARD, DFlash, and DraftTarget workers. ChangesRejection sampling and unified draft sampling refactor
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Init as SpecDecOneEngineForCausalLM
participant Worker as SpecWorkerBase
participant Metadata as SpecMetadata
participant Guard as RejectionBufferGuard
Init->>Worker: set_draft_model(draft_model)
Worker->>Worker: cache self._d2t
Worker->>Metadata: sample_draft_tokens(logits, spec_metadata)
Metadata->>Metadata: prepare_rejection_sampling_buffers()
Worker->>Worker: greedy or advanced sample_draft, scatter draft_probs
Worker->>Guard: validate buffers before acceptance
Guard-->>Worker: valid or invalid
alt buffers invalid
Worker->>Worker: fallback to strict acceptance, clear draft_probs_valid
else buffers valid
Worker->>Worker: compare_and_accept via rejection sampling
end
Possibly related PRs
Suggested reviewers: 🚥 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}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 7
🧹 Nitpick comments (1)
tensorrt_llm/_torch/speculative/interface.py (1)
526-564: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd return annotations to the new helper methods.
Several newly added helpers omit return types, which lets
Anyleak through the central speculative-decoding path. As per coding guidelines, “Always annotate functions. Make the return typeNoneif the function does not return anything.”Also applies to: 1233-1366, 1368-1414, 1574-1799
🤖 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/speculative/interface.py` around lines 526 - 564, The new helper methods in the speculative interface are missing explicit return type annotations, which can leak Any into the decoding path. Add return annotations of None to the relevant no-return helpers, including prepare_rejection_sampling_buffers and prepare, and apply the same rule to the other newly added helper methods in this area. Keep the signatures consistent with the existing class methods in interface.py so the type checker can infer the flow correctly.Source: Coding guidelines
🤖 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.
Inline comments:
In `@tensorrt_llm/_torch/speculative/draft_target.py`:
- Around line 264-271: The final-cycle flag in draft token sampling is using the
configured maximum instead of the actual runtime length, which prevents
finalization when the draft length is shortened. Update the
`sample_draft_tokens()` loop in `DraftTarget` so `is_last_draft_cycle` is
computed from `runtime_draft_len` rather than `self.max_draft_len`, ensuring the
last iteration correctly marks `draft_probs_valid` when the dynamic draft length
is reached.
In `@tensorrt_llm/_torch/speculative/eagle3.py`:
- Around line 874-879: The Eagle3 draft loop in sample_draft_tokens() is using
stale draft probability state because spec_metadata.draft_probs_valid is not
reset before capturing new draft distributions. Update the drafting path around
the runtime_draft_len loop to clear spec_metadata.draft_probs_valid before the
first call to sample_draft_tokens(), matching the other one-model workers, so
each acceptance step only trusts probabilities from the current capture.
In `@tensorrt_llm/_torch/speculative/interface.py`:
- Around line 1249-1261: `maybe_gather_sharded_draft_logits()` currently avoids
gathering for attention-DP batches even when the ADP + LM-head-TP path produces
sharded draft logits, so `advanced_sample_draft()` can sample from only part of
the vocab. Update the sharded-logit detection and gather path so ADP batches
with LM-head TP are included, and make sure the TP gather logic threads
`mapping_lm_head_tp` through instead of falling back to `None`. Adjust the
relevant calls in `maybe_gather_sharded_draft_logits()`,
`advanced_sample_draft()`, and the greedy TP-gather branch to use the correct
mapping for LM-head TP logits.
In `@tensorrt_llm/_torch/speculative/mtp.py`:
- Around line 477-482: The final-draft-step check in `sample_draft_tokens` is
using the full `draft_model.mtp_layers` length, which breaks when
`runtime_draft_len` is shorter and leaves `draft_probs_valid` unset. Update the
loop in `mtp.py` so the `is_last_draft_cycle` argument is based on
`runtime_draft_len` (for example, by comparing `i` against the last captured
draft index) rather than `len(draft_model.mtp_layers) - 1`, ensuring the last
executed draft step is correctly marked for dynamic draft lengths.
In `@tensorrt_llm/_torch/speculative/pard.py`:
- Around line 334-340: The `_reshape_draft_tokens_for_accept` helper currently
reshapes the full `spec_metadata.draft_tokens` buffer to `(0, K)` when `num_gens
== 0`, which can fail for preallocated non-empty buffers. Update the `num_gens
== 0` path in `_reshape_draft_tokens_for_accept` to return an explicit empty
view/slice of `spec_metadata.draft_tokens` with shape `(0, K)` instead of
reshaping the entire buffer, while keeping the existing slicing logic for the
non-empty case.
In `@tensorrt_llm/_torch/speculative/utils.py`:
- Around line 55-68: The config parsing helper in the vocab-size lookup
currently uses a blanket exception handler, which hides unrelated bugs and
triggers the lint warning. In the function that reads config.json and returns
the vocab size, narrow the except block to only the expected failure cases from
file access, JSON parsing, and converting the value to int, while keeping the
existing fallback to target_vocab_size for those cases only.
In `@tests/unittest/_torch/speculative/hw_agnostic/test_dflash.py`:
- Around line 107-134: Add the high-memory test marker to
test_dflash_qwen3_8b_rejection so it is classified with the other CUDA-intensive
speculative rejection tests. Update the test definition in test_dflash.py by
applying the same `@pytest.mark.high_cuda_memory` decoration pattern used
elsewhere in this file or neighboring hw_agnostic speculative tests, keeping the
function name and behavior unchanged.
---
Nitpick comments:
In `@tensorrt_llm/_torch/speculative/interface.py`:
- Around line 526-564: The new helper methods in the speculative interface are
missing explicit return type annotations, which can leak Any into the decoding
path. Add return annotations of None to the relevant no-return helpers,
including prepare_rejection_sampling_buffers and prepare, and apply the same
rule to the other newly added helper methods in this area. Keep the signatures
consistent with the existing class methods in interface.py so the type checker
can infer the flow correctly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 7e95ecf6-d6fe-4685-ba0d-0b4d766b1a5d
📒 Files selected for processing (19)
tensorrt_llm/_torch/model_config.pytensorrt_llm/_torch/models/modeling_speculative.pytensorrt_llm/_torch/models/modeling_utils.pytensorrt_llm/_torch/speculative/SPEC_DEVELOPER_GUIDE.mdtensorrt_llm/_torch/speculative/dflash.pytensorrt_llm/_torch/speculative/draft_target.pytensorrt_llm/_torch/speculative/eagle3.pytensorrt_llm/_torch/speculative/eagle3_dynamic_tree.pytensorrt_llm/_torch/speculative/interface.pytensorrt_llm/_torch/speculative/mtp.pytensorrt_llm/_torch/speculative/pard.pytensorrt_llm/_torch/speculative/utils.pytensorrt_llm/llmapi/llm_args.pytests/integration/test_lists/test-db/l0_h100.ymltests/unittest/_torch/speculative/hw_agnostic/test_dflash.pytests/unittest/_torch/speculative/hw_agnostic/test_draft_target.pytests/unittest/_torch/speculative/hw_agnostic/test_mtp.pytests/unittest/_torch/speculative/hw_agnostic/test_pard.pytests/unittest/_torch/speculative/test_rejection_buffers_guard.py
zhaoyangwang-nvidia
force-pushed
the
rejection-one-model-spec
branch
from
fbd3d2b to
24c911e
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
Author
|
Hi @NVIDIA/trt-llm-llmapi-devs @NVIDIA/trt-llm-torch-devs @NVIDIA/trt-llm-torch-models-devs please help to review this PR, thanks a lot. |
Collaborator
|
PR_Github #57760 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #57760 [ run ] completed with state
|
3 tasks
sunnyqgg
approved these changes
Jul 8, 2026
sunnyqgg
reviewed
Jul 8, 2026
zhaoyangwang-nvidia
force-pushed
the
rejection-one-model-spec
branch
from
0b4b479 to
cb8bac0
Compare
…xt slots Block workers (PARD/DFlash) do not draft context requests (they get a zero placeholder token), leaving their slot-indexed draft_probs rows unwritten. When such a request became a generation request the next iteration, rejection acceptance would read a stale row, which was previously guarded by forcing draft_probs_valid=False for block mixed batches. Fill each context request's slot row with a one-hot placeholder distribution (prob 1.0 at draft-vocab token id 0) at the worker's placeholder-assembly point, so the row is a legal distribution: rejection rejects the placeholder (target_prob(0) ~ 0) and resamples from the target residual, equivalent to strict acceptance. Extracted into SpecWorkerBase.write_context_onehot_draft_probs and shared by PARD and DFlash. This lets block mixed batches keep draft_probs_valid=True (drop the num_contexts==0 clause). Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
With one-model rejection sampling and dynamic draft length, runtime_draft_len is batch-global and can toggle K->0->K as batch size fluctuates. On the 0 iteration every gen request goes through skip_drafting, producing no draft tokens and never scattering draft_probs. The next K iteration then runs rejection acceptance (its branch frozen in the captured CUDA graph, so a Python flag reset cannot redirect it) reading the stale draft_probs rows paired with padded-zero draft tokens, silently corrupting output. Mirror the draft-token zero-padding: mark such gen requests host-side in _handle_dynamic_draft_len (pre-pad, current_num_draft_tokens == 0), collect their slots in _prepare_tp_inputs, and eagerly write a one-hot placeholder distribution (draft-vocab token id 0) into their draft_probs rows after spec_metadata.prepare(). The eager write into the stable draft_probs buffer is read by the replayed rejection kernel, which then rejects the placeholder and resamples from the target -- equivalent to strict acceptance. Idempotent with the context-slot one-hot; no-op when rejection is off. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
The draft_probs_valid producer->consumer flag is now redundant: every slot the rejection acceptance reads is guaranteed to hold a legal distribution written by the request's immediately-preceding step. Context slots in block workers are one-hot'd by write_context_onehot_draft_probs; dynamic-draft-len zero-toggle padding slots by write_padding_onehot_draft_probs; all-greedy and pure-context batches are already excluded by is_all_greedy_sample and the num_gens>0 guard; CUDA-graph warmup only touches the dummy scratch row (real slots are written by each request's own context/gen forward before any rejection reads them). Delete the field, its term in _can_use_rejection_sampling, the fail-closed assignment, reset_draft_probs_valid_for_capture and its five callers, the set in sample_draft_tokens (and the now-dead is_last_draft_cycle parameter), and _invalidate_draft_probs_after_capture. _can_use_rejection_sampling now reduces to `use_rejection_sampling and not is_all_greedy_sample`, matching the dynamic-tree worker's own gate. The _rejection_buffers_valid shape guard is retained. Drop the obsolete warmup-flag test; update the dispatch tests. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
zhaoyangwang-nvidia
force-pushed
the
rejection-one-model-spec
branch
from
84bee1a to
0b59abc
Compare
…n sampling NGram and SA are retrieval-based drafters that emit only draft token ids with no per-token proposal distribution q(x). Rejection sampling needs q to form min(1, p/q) and the (p - q)+ residual, so both are gated. Document the reason on SuffixAutomatonManager and in the validate_speculative_config support whitelist. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #59343 [ run ] triggered by Bot. Commit: |
BowenFu
approved these changes
Jul 15, 2026
Collaborator
|
PR_Github #59343 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #59411 [ run ] triggered by Bot. Commit: |
mzweilz
approved these changes
Jul 15, 2026
Collaborator
|
PR_Github #59411 [ run ] completed with state
|
get_spec_worker holds the correct Mapping but never passed it to the MTP workers: MTPEagleWorker hardcoded super().__init__(mapping=None) and MTPWorker read the unreliable model_config.mapping (None for the MTP worker). With the draft LM head kept vocab-sharded (lm_head_gather_output =False), a None mapping made greedy_sample_draft_with_tp_gather skip the TP argmax gather, so each rank argmaxed its own vocab shard, drafting divergent tokens and desyncing the ranks into a PyExecutor hang under plain TP + overlap scheduler (repro: DeepSeek-V3-Lite fp8 tp4 mtp_nextn=2 cuda_graph+overlap+torch_compile). Thread mapping from get_spec_worker into MTPWorker and MTPEagleWorker, matching every other one-model worker. Verified tp4 output is now byte-identical to the origin/main baseline. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #59450 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #59450 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #59566 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #59566 [ run ] completed with state
|
Under attention DP each rank is data-parallel and owns a distinct set of requests, so the draft logits must not be gathered across ranks -- a per-rank argmax on the rank own logits is the correct proposal. The ADP + LM-head-TP path wrongly classified its vocab-sharded draft logits as needing the TP gather, splicing in an allgather over ranks that hold different token counts. That desynced the ranks and crashed in the DeepEP MoE dispatch (intranode.cu unspecified launch failure) / hung the cuda-graph batch-padding allgather (repro: DeepSeek-V3-Lite tp4 ep4 mtp_nextn=3 attention_dp + enable_lm_head_tp_in_adp, matching CI TestDeepSeekR1 test_nvfp4_multi_gpus[latency_adp_lmtp_tp4]). Make _draft_logits_are_sharded return False for any attention-DP mode so greedy takes a plain argmax and the advanced path skips the gather, matching the pre-rework behavior. Remove the now-unreachable ADP+LM-head-TP gather branches. Verified: ADP+LMTP tp4 ep4 and plain-TP tp4 both produce correct output byte-matching the baseline. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #59593 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #59593 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
yingguo-trt
approved these changes
Jul 16, 2026
Collaborator
|
PR_Github #59652 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #59652 [ run ] completed with state |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
10 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
New Features
Bug Fixes
Documentation
Tests
Description
Extends rejection sampling (previously available only for Eagle3 dynamic-tree)
to the one-model speculative-decoding workers: vanilla MTP, PARD, DFlash, and
DraftTarget one-model. Rejection sampling makes speculative decoding
distribution-preserving under non-greedy sampling (temperature / top-k / top-p),
so accepted tokens follow the target model's distribution instead of falling back
to strict token-equality acceptance.
Alongside the feature, this PR consolidates the draft-token production and
acceptance paths that had diverged across the five one-model workers into a
single shared implementation on
SpecWorkerBase, so each method now differs onlywhere the model architecture forces it to.
What changed
Feature: rejection sampling for one-model methods
slot-indexed
draft_probsbuffer (keyed bypy_seq_slot), consumed by therejection kernel on the next forward.
draft_probs_validis reset at the start of everydraft forward and only trusted after a complete capture;
_rejection_buffers_validvalidates every dereferenced shape (CUDA-graph-safe, no host sync) and falls
back to strict acceptance on any mismatch.
llm_args.py: rejection is opt-in per method; unsupportedcombinations (SA, relaxed-thinking, attention-DP, context-parallel, guided
decoding, missing flashinfer) either raise on explicit opt-in or silently
disable by default. Plain TP is supported via minimal draft-logit gather.
Refactor: unified production / acceptance
sample_draft_tokens(step form for MTP/DraftTarget,block form for PARD/DFlash); Eagle3 now routes through it too. d2t is cached once
at draft-model load (
self._d2t) and applied internally — no longer threadedthrough call arguments.
sample_and_accept_draft_tokens->compare_and_accept;workers override only the
_reshape_draft_tokens_for_accept/_reshape_logits_for_accepthooks when their buffer/logit layout differs (onlyPARD does). MTP and Eagle3 keep full overrides for their relaxed / THOP paths.
(
lm_head_gather_output=False); the base gathers the minimum necessary based ona shape check (
_draft_logits_are_sharded) — light index+value gather for greedy,full all-gather only for a sharded non-greedy batch.
Docs & tests
tensorrt_llm/_torch/speculative/SPEC_DEVELOPER_GUIDE.md, anagent-friendly module guide (so both humans and coding agents can extend the
speculative code without re-deriving the framework). It documents the
SpecWorkerBase inheritance framework, d2t handling, the do-not-gather
draft-logit policy, the acceptance hooks, and the rejection-sampling contract a
new worker must satisfy. It follows the same in-tree developer-guide convention
as
tensorrt_llm/_torch/modules/ATTENTION_DEVELOPER_GUIDE.mdandtensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md(and, forcriteria-style docs,
tensorrt_llm/_torch/visual_gen/ENGINEERING_CRITERIA.md) —these are already referenced from
AGENTS.mdas read-before-modifying guides.(test_draft_target, test_pard, test_dflash, test_mtp); plus
test_rejection_buffers_guard for the buffer-allocation and fail-closed guard
logic. Registered in the H100 pre-merge stage.
Test Coverage
tests/unittest/_torch/speculative/hw_agnostic/test_{draft_target,pard,dflash,mtp}.py— one rejection-on non-greedy e2e per method (path-correctness).
tests/unittest/_torch/speculative/test_rejection_buffers_guard.py— bufferallocation + fail-closed guard + acceptance dispatch (parametrized).
The draft_target / pard / dflash / eagle3 matrix (rejection on/off × cuda-graph)
matches the pre-refactor baseline; MTP's pre-existing REJ_OFF+GRAPH_ON crash is
an unrelated upstream issue.
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-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin 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.