[TRTLLM-11508][refactor] Merge Eagle3 and MTP-eagle one-model workers by zhaoyangwang-nvidia · Pull Request #12353 · NVIDIA/TensorRT-LLM

zhaoyangwang-nvidia · 2026-03-19T09:24:15Z

Summary by CodeRabbit

Release Notes

New Features
- Added support for MTP Eagle speculative decoding mode for enhanced inference acceleration alongside existing Eagle3 support.
- Introduced relaxed acceptance configuration during reasoning/thinking phases to optimize token generation.
- Extended speculative decoding compatibility to additional model architectures (DeepseekV3, ExaoneMoe, GLM, Nemotron, Qwen3).

Description

Unify Eagle3OneModelWorker (Eagle3 one-model) and MTPEagleWorker
(MTP-eagle one-model) into a single worker class in
tensorrt_llm/_torch/speculative/eagle3.py, branching on
self.is_mtp_eagle = spec_dec_mode.is_mtp_eagle_one_model(). The two
code paths were ~85% duplicated; this PR collapses them into one
implementation while preserving backward-compatible imports.

MTPEagleWorker becomes a thin backward-compatible subclass in
eagle3.py. mtp.py retains a module-level __getattr__ shim so
from tensorrt_llm._torch.speculative.mtp import MTPEagleWorker
continues to resolve.

Key changes

Unified worker (eagle3.py): new helpers
_get_step_all_rank_num_tokens, _run_draft_forward,
_prepare_flash_mla_generation_layout, draft_sampler (TP-aware);
sample_and_accept_draft_tokens gains an input_ids parameter and
the relaxed-thinking acceptance path previously exclusive to
MTPEagleWorker.
Unified metadata (Eagle3OneModelSpecMetadata): new fields
slot_ids and subseq_all_rank_num_tokens. prepare() skips
num_tokens subtraction for MTP-eagle and populates slot_ids from
the resource manager.
Eagle3ResourceManager owns the relaxed-acceptance
relaxed_delta_pool so the Eagle3 path can also use thinking-phase
relaxed acceptance.
Eagle3DraftModel.forward takes an optional
all_rank_num_tokens kwarg and wraps its body in try/finally that
restores attn_metadata.all_rank_num_tokens on exit — the worker
loop no longer mutates attn_metadata for Eagle3.
EagleDecodingConfig gains the five relaxed-acceptance fields
mirrored from MTPDecodingConfig.
SpeculativeDecodingMode.is_mtp_one_model() is narrowed to
vanilla MTP only; MTP_EAGLE_ONE_MODEL becomes a first-class
one-model mode in use_one_engine, without_logits,
needs_kv_cache_rewind, support_overlap_scheduler,
support_capturable_guided_decoder, support_dynamic_draft_len,
has_spec_decoder. Per-model checks in modeling_deepseekv3,
modeling_glm, modeling_exaone_moe, modeling_nemotron_h,
modeling_qwen3_next, modeling_speculative, and model_config
are extended accordingly.
Factory routing (utils.py) routes MTP_EAGLE_ONE_MODEL to
Eagle3OneModelSpecMetadata, Eagle3OneModelSampler,
Eagle3ResourceManager, and the unified worker.
model_engine.py populates
spec_metadata.subseq_all_rank_num_tokens for both Eagle3 and
MTP-eagle one-model at all three attention-DP allgather sites.

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)
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.

zhaoyangwang-nvidia · 2026-03-19T10:31:09Z

/bot run

tensorrt-cicd · 2026-03-19T10:36:56Z

PR_Github #39589 [ run ] triggered by Bot. Commit: 9f923d9 Link to invocation

tensorrt-cicd · 2026-03-19T14:01:55Z

PR_Github #39589 [ run ] completed with state SUCCESS. Commit: 9f923d9
/LLM/main/L0_MergeRequest_PR pipeline #30802 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

Link to invocation

zhaoyangwang-nvidia · 2026-03-20T10:16:35Z

/bot run

tensorrt-cicd · 2026-03-20T10:23:07Z

PR_Github #39733 [ run ] triggered by Bot. Commit: 01b1fe7 Link to invocation

tensorrt-cicd · 2026-03-20T14:00:19Z

PR_Github #39733 [ run ] completed with state SUCCESS. Commit: 01b1fe7
/LLM/main/L0_MergeRequest_PR pipeline #30928 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

Link to invocation

zhaoyangwang-nvidia · 2026-03-22T08:23:13Z

/bot run

tensorrt-cicd · 2026-03-22T08:28:48Z

PR_Github #39814 [ run ] triggered by Bot. Commit: 5a1ef63 Link to invocation

tensorrt-cicd · 2026-03-22T10:21:48Z

PR_Github #39814 [ run ] completed with state SUCCESS. Commit: 5a1ef63
/LLM/main/L0_MergeRequest_PR pipeline #30991 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

Link to invocation

zhaoyangwang-nvidia · 2026-03-22T10:39:43Z

/bot run

tensorrt-cicd · 2026-03-22T10:46:34Z

PR_Github #39819 [ run ] triggered by Bot. Commit: 5a1ef63 Link to invocation

tensorrt-cicd · 2026-03-22T12:43:21Z

PR_Github #39819 [ run ] completed with state SUCCESS. Commit: 5a1ef63
/LLM/main/L0_MergeRequest_PR pipeline #30996 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

Link to invocation

zhaoyangwang-nvidia · 2026-03-22T14:03:10Z

/bot run

tensorrt-cicd · 2026-03-22T14:10:00Z

PR_Github #39828 [ run ] triggered by Bot. Commit: 5a1ef63 Link to invocation

tensorrt-cicd · 2026-03-22T16:06:18Z

PR_Github #39828 [ run ] completed with state SUCCESS. Commit: 5a1ef63
/LLM/main/L0_MergeRequest_PR pipeline #31005 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

Link to invocation

zhaoyangwang-nvidia · 2026-03-26T07:05:07Z

/bot run

tensorrt-cicd · 2026-03-26T07:10:40Z

PR_Github #40431 [ run ] triggered by Bot. Commit: f4f8ac6 Link to invocation

tensorrt-cicd · 2026-03-26T11:13:45Z

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

Link to invocation

lfr-0531

LGTM~

tensorrt-cicd · 2026-06-02T12:38:17Z

PR_Github #51578 [ run ] completed with state SUCCESS. Commit: 9051c11
/LLM/main/L0_MergeRequest_PR pipeline #40968 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

zhaoyangwang-nvidia · 2026-06-02T12:45:53Z

/bot run

zhaoyangwang-nvidia · 2026-06-03T01:36:18Z

/bot run

tensorrt-cicd · 2026-06-03T01:41:51Z

PR_Github #51718 [ run ] triggered by Bot. Commit: 9051c11 Link to invocation

tensorrt-cicd · 2026-06-03T06:20:45Z

PR_Github #51718 [ run ] completed with state SUCCESS. Commit: 9051c11
/LLM/main/L0_MergeRequest_PR pipeline #41094 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

zhaoyangwang-nvidia · 2026-06-03T06:33:35Z

/bot run

tensorrt-cicd · 2026-06-03T06:39:52Z

PR_Github #51777 [ run ] triggered by Bot. Commit: 9051c11 Link to invocation

tensorrt-cicd · 2026-06-03T11:03:33Z

PR_Github #51777 [ run ] completed with state SUCCESS. Commit: 9051c11
/LLM/main/L0_MergeRequest_PR pipeline #41143 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

zhaoyangwang-nvidia · 2026-06-04T01:45:50Z

/bot run

tensorrt-cicd · 2026-06-04T01:52:56Z

PR_Github #51934 [ run ] triggered by Bot. Commit: 9051c11 Link to invocation

tensorrt-cicd · 2026-06-04T02:50:17Z

PR_Github #51934 [ run ] completed with state SUCCESS. Commit: 9051c11
/LLM/main/L0_MergeRequest_PR pipeline #41286 completed with status: 'SUCCESS'

CI Report

Link to invocation

…NVIDIA#12353) Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com> Signed-off-by: NVFB <186336021+NVFB@users.noreply.github.com>

…NVIDIA#12353) Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>

…revert Remove all 14 deepseek-v32-fp4 perf-sanity test waivers (nvbugs 6280721, 6280649, 6085022, 6200257). These were skipped due to the gen-server CUDA-graph capture crash introduced by NVIDIA#12353, which is reverted in this branch, so the cases can run again.

…DIA#12353)" This reverts commit 8e5d9e2 (NVIDIA#12353). The MTP-eagle / Eagle3 one-model worker merge regresses the GB200 2-node disagg perf-sanity gen-only test: test_perf_sanity.py::test_e2e[disagg_upload-gen_only- gb200_deepseek-v32-fp4_1k1k_con2048_ctx1_dep4_gen1_dep4_eplb0_mtp1_ccb-NIXL] The GEN server (MTP speculative decoding, num_nextn_predict_layers=1) crashes during executor init while capturing the generation-phase CUDA graphs: _run_cuda_graph_warmup -> _capture_generation_cuda_graphs -> torch.cuda.synchronize() torch.AcceleratorError: CUDA error: unspecified launch failure All 4 gen ranks abort (exit 134). Bisect between the last passing post-merge commit (33b0a32) and the first failing one (316430f) isolates NVIDIA#12353 as the only change touching the spec-decode / model_engine CUDA-graph capture path; the remaining commits in the range are AutoDeploy, test, and infra changes. Reverting restores the gen server startup and unblocks the perf-sanity test. Signed-off-by: Chenfei Zhang <chenfeiz@nvidia.com>

…revert Remove all 14 deepseek-v32-fp4 perf-sanity test waivers (nvbugs 6280721, 6280649, 6085022, 6200257). These were skipped due to the gen-server CUDA-graph capture crash introduced by NVIDIA#12353, which is reverted in this branch, so the cases can run again. Signed-off-by: Chenfei Zhang <chenfeiz@nvidia.com>

Resolve conflicts in speculative decoding introduced by upstream's "[TRTLLM-11508][refactor] Merge Eagle3 and MTP-eagle one-model workers" (NVIDIA#12353): - interface.py support_dynamic_draft_len(): take the union of both sides' mode lists (this PR's is_mtp_one_model/is_pard/is_dflash/ is_draft_target_one_model/is_sa plus upstream's is_mtp_eagle_one_model). - mtp.py: drop the local MTPEagleWorker. Upstream moved it to eagle3.py as a thin subclass of Eagle3OneModelWorker, which already implements the dynamic draft-length behavior this PR added; __init__.py/utils.py now import it from eagle3.py. - mtp.py change_attn_metadata(): upstream added a kv_lens_cuda clamp using its local mtp_num_modules, which this PR had renamed to runtime_draft_len; the auto-merge left a dangling reference, fixed to clamp(min=runtime_draft_len). Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>

…-model merge #12353) Squashed probe equal to the tree of upstream commit 8e5d9e2 ([TRTLLM-11508][refactor] Merge Eagle3 and MTP-eagle one-model workers #12353), rebased onto its parent 910826b so CI builds exactly that commit's tree. Authored/signed-off by chenfeiz to satisfy DCO (original commit's signoff name did not match its author). Signed-off-by: chenfeiz <chenfeiz@nvidia.com>

Qwen3NextMTP.forward accepted all_rank_num_tokens and dropped it with `del`, so the MoE inside a draft iteration kept reading the target forward's per-rank token counts off attn_metadata. After step 0 an MTP Eagle draft processes one token per sequence, so those counts are stale: the scheduler derives the chunk count from them and AllGatherReduceScatter derives the allgather/reducescatter `sizes` from them. The declared sizes then disagree with the actual buffers, and because every rank shares the same wrong list the failure is a collective stall rather than an error. Install the draft distribution for the duration of the call and restore the target value in `finally`, so an exception cannot leak the draft counts into later forwards. Passing None leaves the metadata untouched, which keeps every non-ADP and non-draft caller unchanged. This closes a leg left open by NVIDIA#12353 ("Merge Eagle3 and MTP-eagle one-model workers"), which moved save/restore from the caller to the callee: the Eagle3 leg gained the try/finally in modeling_speculative.py, but the MTP Eagle leg lands in modeling_qwen3_next.py, which that PR did not touch. The block added here has the same shape as modeling_speculative.py:426 and :681. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>

Qwen3NextMTP.forward accepted all_rank_num_tokens and dropped it with `del`, so the MoE inside a draft iteration kept reading the target forward's per-rank token counts off attn_metadata. The draft loop rewrites the sequence layout in place after the first step (eagle3.py:1110 fills seq_lens with 1), so from step 1 the tensor carries one row per sequence while the counts still describe the target's max_draft_len + 1 rows per sequence. Step 0 is unaffected. spec_metadata.num_tokens is deliberately not reduced to the subseq shape for MTP Eagle (eagle3.py:519, "keep the 1st-iter shape (matches input_ids)"), so the value already on attn_metadata describes that step. With max_draft_len == 1 there is no divergence at all. The consumer decides what the staleness costs. AllGatherReduceScatter passes the list straight to allgather as exact per-rank sizes, and _allgather asserts input.shape[dim] == sizes[rank] (ops.py:247) before any NCCL call, so every rank raises AssertionError during dispatch: attention DP plus MTP Eagle with max_draft_len > 1 cannot run on that strategy at all. That combination is not exotic, because TRTLLM_FORCE_COMM_METHOD=ALLGATHER is the standard workaround for DeepEP and NVLink one-sided issues, so the fallback is broken exactly when it is needed. The alltoall strategies instead treat the value as an upper bound: reducescatter is skipped entirely (interface.py:1175) and NVLinkOneSided takes only max() for workspace capacity, so those paths already computed the correct result. What they pay is wasted work. calculate_num_chunks derives the chunk count from len(list) * max(list), so a stale value inflates it by up to max_draft_len + 1, and each surplus chunk is substituted with chunk 0's tokens (moe_scheduler.py:659) and run through a full dispatch, expert GEMM and combine before its result is discarded. The saving only materialises once the padded row count crosses moe_max_num_tokens, so it is configuration dependent and is not measured here. On SM120, where use_dp_padding is enabled, the same stale value truncates the output (moe_scheduler.py:216) and returns the target row count instead of the draft's. No numerical result changes on any validated deployment. Install the draft distribution for the duration of the call and restore the target value in `finally`, so an exception cannot leak the draft counts into later forwards. Passing None leaves the metadata untouched, which keeps every non-ADP and non-draft caller unchanged. This closes a leg left open by NVIDIA#12353 ("Merge Eagle3 and MTP-eagle one-model workers"), which moved save/restore from the caller to the callee: the Eagle3 leg gained the try/finally in modeling_speculative.py, but the MTP Eagle leg lands in modeling_qwen3_next.py, which that PR did not touch. The block added here has the same shape as modeling_speculative.py:426 and :681. test_qwen3_next_eager_fusion.py arrived with NVIDIA#15194 and has never appeared in a test list, so the cases added here would not have run in CI. Register the file in l0_a10.yml next to test_qwen3_next_moe_quant.py, which the comment above that entry already marks as the parking spot for CPU-only unit tests with no dedicated job. The CPU stage is not an option: it runs `-m cpu_only`, and conftest.py:228 skips any test file carrying no pytest.mark.cpu_only marker, which fails the entry with exit code 5 for collecting nothing. Marking the file is not the alternative it looks like. L0_Test.groovy:1578 hands every stage whose name does not start with CPU- the opposite expression, `-m 'not cpu_only'`, so a marked file would deselect all six cases on l0_a10 and fail there in exactly the same way. The two options are mutually exclusive: either the file carries the marker and is listed on the CPU stage, or it stays unmarked and is listed on a GPU stage. This takes the second. Both directions were checked against the file: `-m 'not cpu_only'` collects and passes all six, `-m cpu_only` deselects all six and exits 5. All six, including the four that predate this change, also pass with CUDA_VISIBLE_DEVICES empty. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>

github-actions Bot assigned zhaoyangwang-nvidia Mar 19, 2026

zhaoyangwang-nvidia force-pushed the merge-eagle-mtp branch from 98b1417 to 9f923d9 Compare March 19, 2026 09:52

zhaoyangwang-nvidia force-pushed the merge-eagle-mtp branch from 9f923d9 to 01b1fe7 Compare March 20, 2026 10:15

zhaoyangwang-nvidia force-pushed the merge-eagle-mtp branch from 01b1fe7 to 5a1ef63 Compare March 22, 2026 08:23

zhaoyangwang-nvidia force-pushed the merge-eagle-mtp branch 2 times, most recently from 7e82c9c to f4f8ac6 Compare March 24, 2026 09:06

zhaoyangwang-nvidia changed the title ~~[TRTLLM-11508][refactor] Merge eagle mtp~~ [TRTLLM-11508][refactor] Merge Eagle3 and MTP-eagle one-model workers May 11, 2026

zhaoyangwang-nvidia force-pushed the merge-eagle-mtp branch 2 times, most recently from 8ed37a4 to 67e4428 Compare May 12, 2026 09:57

zhaoyangwang-nvidia marked this pull request as ready for review May 12, 2026 09:58

zhaoyangwang-nvidia requested review from a team as code owners May 12, 2026 09:58

lfr-0531 approved these changes Jun 2, 2026

View reviewed changes

mikeiovine merged commit 8e5d9e2 into NVIDIA:main Jun 4, 2026
7 checks passed

2ez4bz pushed a commit to 2ez4bz/TensorRT-LLM that referenced this pull request Jun 8, 2026

[TRTLLM-11508][refactor] Merge Eagle3 and MTP-eagle one-model workers (…

e931bbc

…NVIDIA#12353) Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>

tensorrt-cicd mentioned this pull request Jun 10, 2026

[None][fix] Revert "Merge Eagle3 and MTP-eagle one-model workers (#12353)" #15217

Closed

tensorrt-cicd mentioned this pull request Jun 10, 2026

[None][test] CI-only: RUN (not skip) gb200 deepseek-v32 gen-only perf-sanity at 33b0a32 (DO NOT MERGE) #15224

Closed

This was referenced Jun 11, 2026

[None][test] CI verify L0_PostMerge 2767 (b4d44d3): disagg-v32-fp4 NIXL gen_only IMA #15241

Merged

[None][test] CI bisect nvbugs/6280721: probe efb71c7 / pm-2755 (Cherry-pick kv_cache_manager_v2 fixes to main #14725) #15242

Closed

zhaoyangwang-nvidia deleted the merge-eagle-mtp branch July 31, 2026 09:50

Uh oh!

Conversation

zhaoyangwang-nvidia commented Mar 19, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Release Notes

Description

Key changes

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

zhaoyangwang-nvidia commented Mar 19, 2026

Uh oh!

tensorrt-cicd commented Mar 19, 2026

Uh oh!

tensorrt-cicd commented Mar 19, 2026

Uh oh!

zhaoyangwang-nvidia commented Mar 20, 2026

Uh oh!

tensorrt-cicd commented Mar 20, 2026

Uh oh!

tensorrt-cicd commented Mar 20, 2026

Uh oh!

zhaoyangwang-nvidia commented Mar 22, 2026

Uh oh!

tensorrt-cicd commented Mar 22, 2026

Uh oh!

tensorrt-cicd commented Mar 22, 2026

Uh oh!

zhaoyangwang-nvidia commented Mar 22, 2026

Uh oh!

tensorrt-cicd commented Mar 22, 2026

Uh oh!

tensorrt-cicd commented Mar 22, 2026

Uh oh!

zhaoyangwang-nvidia commented Mar 22, 2026

Uh oh!

tensorrt-cicd commented Mar 22, 2026

Uh oh!

tensorrt-cicd commented Mar 22, 2026

Uh oh!

zhaoyangwang-nvidia commented Mar 26, 2026

Uh oh!

tensorrt-cicd commented Mar 26, 2026

Uh oh!

tensorrt-cicd commented Mar 26, 2026

Uh oh!

lfr-0531 left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented Jun 2, 2026

Uh oh!

zhaoyangwang-nvidia commented Jun 2, 2026

Uh oh!

zhaoyangwang-nvidia commented Jun 3, 2026

Uh oh!

tensorrt-cicd commented Jun 3, 2026

Uh oh!

tensorrt-cicd commented Jun 3, 2026

Uh oh!

zhaoyangwang-nvidia commented Jun 3, 2026

Uh oh!

tensorrt-cicd commented Jun 3, 2026

Uh oh!

tensorrt-cicd commented Jun 3, 2026

Uh oh!

zhaoyangwang-nvidia commented Jun 4, 2026

Uh oh!

tensorrt-cicd commented Jun 4, 2026

Uh oh!

tensorrt-cicd commented Jun 4, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

zhaoyangwang-nvidia commented Mar 19, 2026 •

edited by coderabbitai Bot

Loading