[None][feat] Support rejection sampling under attention DP (incl. LM-head TP)#16544
Open
zhaoyangwang-nvidia wants to merge 4 commits into
Open
Conversation
…cl. LM-head TP Lift the attention-DP gate on one-model rejection sampling. Three pieces: 1. Group-synchronize is_all_greedy_sample under ADP+LM-head-TP+rejection: the greedy-vs-advanced choice gates the LM-head-TP group's collectives (stacked draft forward, distributed argmax) and the CUDA graph variant recording them, but each ADP rank owns a different batch so local flags can diverge. The model engine all-gathers the local flags before the CUDA graph key is built and stores the group AND on spec_metadata; _scan_one_model_sampling re-applies it on every rescan so the graph key, buffer population, and worker branches all agree. AND semantics are conservative-safe: greedy ranks pulled onto the advanced path sample greedily via their sentinel params. 2. Bypass the LM-head-TP fast path for advanced draft sampling: rejection needs full per-request distributions, which the stacked/vocab-sharded layout cannot provide without a full-logits gather. Under ADP the lm_head weight is replicated, so the worker calls lm_head directly (plain local full-vocab GEMM, same as the target head) whenever rejection is enabled and the (group-synced) batch is not all-greedy. With rejection off, non-greedy batches keep the LM-head-TP argmax path unconditionally and the per-rank flag never enters control flow. 3. llm_args: drop enable_attention_dp from the rejection parallel gate (context parallelism remains gated); update messages and stale comments. Plain-ADP ranks hold replicated full-vocab draft logits, so their advanced path is local-only and needs no synchronization. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
…cument sync trade-offs Review cleanups: extract the greedy-vs-advanced decision into SpecMetadata.wants_advanced_draft_sampling so the sampler branch and the worker's LM-head-TP bypass cannot diverge; drop getattr defaults that defended fields the SpecMetadata dataclass always declares; document why the group flag sync is a dedicated all-gather (the all_rank_num_tokens exchange runs after the CUDA graph key is built) and why the live flag is overwritten in addition to storing the rescan override. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
zhaoyangwang-nvidia
added
the
api-compatible
…cy baseline D205/D209 in the new test_llm_args gate tests; regenerate ruff-legacy-baseline.json (this branch resolves 4 legacy violations and the regeneration drops stale entries for files no longer in the tree). Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
zhaoyangwang-nvidia
marked this pull request as ready for review
July 17, 2026 10:09
zhaoyangwang-nvidia
requested review from
Tabrizian,
cascade812,
mikeiovine,
mlefeb01,
reasonsolo and
tburt-nv
July 17, 2026 10:09
- Drop conditions defending unreachable states: Mapping.__init__ asserts enable_lm_head_tp_in_adp implies enable_attention_dp, so the separate ADP check and the getattr default are redundant (invariant cited in comments). - Single-source the group-sync semantics: SpecMetadata.group_all_greedy_sample is the anchor; the model-engine sync, the scan application, and the worker branch keep only their local contract plus a pointer. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (7)
📝 WalkthroughWalkthroughThe change synchronizes speculative decoding decisions across tensor-parallel ranks, permits rejection sampling with attention parallelism, adjusts MTP-Eagle advanced-sampling routing, adds unit coverage, and refreshes Ruff legacy baseline totals and entries. ChangesSpeculative sampling synchronization
Ruff baseline refresh
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant PyTorchModelEngine
participant SpecMetadata
participant TensorParallelGroup
participant CUDAGraphKeySelection
PyTorchModelEngine->>SpecMetadata: update_is_all_greedy_sample
PyTorchModelEngine->>TensorParallelGroup: all-gather local greedy flags
TensorParallelGroup-->>PyTorchModelEngine: group-wide AND result
PyTorchModelEngine->>SpecMetadata: store synchronized sampling state
PyTorchModelEngine->>CUDAGraphKeySelection: construct sampling-path key
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
mikeiovine
approved these changes
Jul 17, 2026
Comment on lines
+519
to
+533
| # Under ADP + LM-head TP with rejection sampling, the greedy-vs-advanced | ||
| # sampling-path choice gates group collectives (the LM-head-TP stacked | ||
| # draft forward and its distributed argmax) and the CUDA graph variant | ||
| # recording them. Under attention DP each rank owns a different batch, so | ||
| # the local flags can diverge -- a diverged choice would leave part of the | ||
| # group waiting in a collective the rest never enters. The model engine | ||
| # (``_sync_group_all_greedy_sample``) therefore all-gathers the per-rank | ||
| # local flags each iteration BEFORE the CUDA graph key is built and stores | ||
| # the group AND here; ``_scan_one_model_sampling`` re-applies it on every | ||
| # rescan so all later consumers (graph key, populate, worker branches) | ||
| # observe the same value. AND semantics are conservative-safe: a locally | ||
| # all-greedy rank may be pulled onto the advanced path, where its sentinel | ||
| # sampling params still sample greedily -- never the reverse. ``None`` | ||
| # (default) means "no group sync configured": the local per-rank value is | ||
| # used as-is. |
Collaborator
There was a problem hiding this comment.
This comment is very verbose when the behavior is pretty simple and easy to understand - consider simplifying it
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
2 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
Bug Fixes
Tests
Description
Rejection sampling (introduced in #15775; makes speculative decoding
distribution-preserving under non-greedy sampling) was config-gated off under
attention DP. This PR lifts that gate and makes the advanced (rejection) draft
path work correctly under ADP, including the
enable_lm_head_tp_in_adpmode.Why it was gated
Under ADP + LM-head TP, the MTP draft head produces group-stacked,
vocab-sharded logits (
[lm_head_tp_size * max_num_requests, vocab/g]). Greedydrafting digests this layout with a cheap distributed argmax (#16440), but
rejection sampling needs each request's full-vocab proposal distribution
q(x) to compute
min(1, p/q)and the(p - q)+residual, which vocab shardscannot provide. Additionally, the greedy-vs-advanced branch is driven by
is_all_greedy_sample— a per-rank batch property under ADP — so a divergedchoice would strand part of the LM-head-TP group in a collective the rest
never enters (hang), or replay CUDA graphs with mismatched collective
sequences.
Note: the parallel gate only ever covered the newly wired methods (vanilla
MTP / PARD / DFlash / DraftTarget one-model); MTP-Eagle was never gated, so on
current main MTP-Eagle + ADP + LM-head-TP + rejection + a non-greedy batch can
already reach the advanced path and silently sample from the stacked/sharded
layout. For MTP-Eagle this PR is a correctness fix; for the other methods it
is new support.
What changed
Group-synchronized
is_all_greedy_sample(interface.py,model_engine.py): under ADP + LM-head-TP + rejection, the model engineall-gathers the per-rank flags BEFORE the CUDA graph key is built and
stores the group AND on
spec_metadata.group_all_greedy_sample;_scan_one_model_samplingre-applies it on every rescan so the graph key,buffer population, and worker branches all observe the same value. AND
semantics are conservative-safe: a locally all-greedy rank pulled onto the
advanced path still samples greedily via its sentinel params. The sync
condition is pure config, so ranks also agree on whether the exchange
happens at all.
LM-head-TP bypass for advanced draft sampling (
eagle3.py): whenwants_advanced_draft_sampling(rejection enabled AND group-synced batchnot all-greedy), the worker skips the stacked/sharded fast path and calls
draft_model.lm_headdirectly — under ADP the LMHead weight is replicated(the shard slicing is a runtime optimization behind
is_spec_decoding_head=True), so this is a plain local full-vocab GEMM,identical to what the target head runs, with zero extra communication.
With rejection off, non-greedy batches keep the LM-head-TP argmax path
unconditionally and the per-rank flag never enters control flow.
Lift the config gate (
llm_args.py): removeenable_attention_dpfrom the rejection parallel gate (context parallelism remains gated).
Plain ADP needs no synchronization: each rank holds replicated full-vocab
draft logits for its own requests and the advanced path is local-only.
The greedy-vs-advanced predicate is centralized in
SpecMetadata.wants_advanced_draft_samplingso the sampler branch and theworker's head selection cannot diverge.
Test Coverage
Unit (CPU):
tests/unittest/_torch/speculative/test_group_all_greedy_sync.py(new, 5 cases: local vs group override, rescan survival, capture-override
composition);
tests/unittest/llmapi/test_llm_args.py(2 new gate cases: ADPincl. LM-head-TP accepted, context parallelism still raises);
test_rejection_buffers_guard.pypasses unchanged.E2E (4xH200, DeepSeek-V3-Lite bf16, tp4 + attention DP + LM-head TP, MTP
nextn=2, rejection enabled, CUDA graphs on):
Related
use_lm_head_tp_in_adp and is_all_greedy_sample, keep… #16440 (greedy distributed argmax for ADP + LM-head TP, nvbug 6460072)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.Note on
ruff-legacy-baseline.jsonThis PR resolves 4 legacy lint violations in the files it touches, so the
ruff-legacybaseline was tightened per the hook's suggestion (scripts/legacy_utils.py lint-update-violations). The regeneration is a full rescan and also drops stale baseline entries for files that no longer exist in the tree — that part is mechanical cleanup with no behavior change.