[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

Eagle3OneModelWorker on the MTP Eagle path consumes the target model's hidden_states argument directly (prepare_1st_drafter_inputs and _run_draft_forward both gate the spec_metadata.hidden_states read on not is_mtp_eagle), so the layer-capture machinery is never read. Configuring layers_to_capture={last} for MTP_EAGLE_ONE_MODEL still had two visible side effects: - modeling_deepseekv3 / modeling_glm / modeling_gpt_oss disable POST_MOE_FUSION / POST_MLP_FUSION whenever spec_metadata.is_layer_capture( layer_idx) is true. MTP Eagle paid that cost despite never using the capture. - Eagle3OneModelSpecMetadata allocated a (max_num_tokens, hidden_size) hidden_states buffer (and Eagle3ResourceManager pre-allocated one too) that the MTP Eagle worker never reads. Opt out of capture for MTP_EAGLE_ONE_MODEL: - MTPDecodingConfig.num_capture_layers returns 0 for one-model (1 stays for two-model MTP_EAGLE, which still feeds captured hidden states into the separate draft engine). - Eagle3OneModelSpecMetadata.__post_init__ defaults layers_to_capture to () for MTP_EAGLE_ONE_MODEL and skips the hidden_states allocation. - get_spec_metadata() stops passing layers_to_capture={last} on the MTP_EAGLE_ONE_MODEL branch and lets the metadata default kick in. With layers_to_capture empty, is_layer_capture() is false for every layer and maybe_capture_hidden_states() is a no-op for-loop, so post-MLP / post-MoE fusion stays enabled and no capture buffer is allocated. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>

Bundle of small fixes from the merge-eagle-mtp PR review: - Drop the historical `kv_lens_cuda` save/restore NOTE in `Eagle3OneModelWorker._prepare_attn_metadata_for_spec_dec()`. The note was internal context for the merge commit and is not useful as inline documentation of the current state. - Remove the `if self.is_mtp_eagle:` gate around the Mamba-hybrid state update in `Eagle3OneModelWorker.forward()`. The `isinstance(kv_cache_manager, MambaHybridCacheManager)` check already no-ops on non-Mamba caches, so spec-mode gating was unnecessary and would prevent Eagle3 from picking up Mamba support if a Mamba-style Eagle draft ever lands. - Delete the `__getattr__` backward-compat shim in `mtp.py` that re-exported `MTPEagleWorker` from `eagle3.py`. No in-tree caller imports `MTPEagleWorker` from `tensorrt_llm._torch.speculative.mtp` anymore (both `speculative/__init__.py` and `speculative/utils.py` import it directly from `.eagle3`). - Drop the relaxed-acceptance fields (`use_relaxed_acceptance_for_thinking`, `relaxed_topk`, `relaxed_delta`, `begin_thinking_phase_token`, `end_thinking_phase_token`) from `EagleDecodingConfig`. The unified worker's relaxed-acceptance branch is already guarded by `getattr(spec_config, 'use_relaxed_acceptance_for_thinking', False)`, so an `EagleDecodingConfig` without these fields naturally short- circuits the path. `MTPDecodingConfig` keeps the same fields unchanged. Limits the user-facing surface of relaxed acceptance to the MTP family, per review preference (relaxed acceptance degrades the spec-dec output distribution and is questionable for Eagle). - Rewrite the "Eagle engine takes ``draft_len`` tokens" comment in `_forward_linear_draft_loop()` to mention both Eagle3 and MTP Eagle, matching the unified worker that path now drives. - Expand the comment around the `num_tokens` subtract in `Eagle3OneModelSpecMetadata.prepare()` to explain that the value is consumed only by the attention-DP allgather in `model_engine`, that both modes share `prepare_1st_drafter_inputs()` and therefore feed the same input_ids tensor to the kernel, and that the subtract asymmetry exists solely to preserve the pre-refactor shape-hint convention of each mode. The previous wording ("MTP Eagle uses max_draft_len + 1 tokens... Eagle3 follows the standard tree/linear path") incorrectly implied an input-shape difference. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>

Factor the duplicated spec_metadata distributed-token-count assignments (all_rank_num_tokens, all_rank_num_seqs, and the Eagle3/MTP-eagle one-model subseq_all_rank_num_tokens) into a shared helper _set_spec_metadata_all_rank_num_tokens, called from _apply_incremental_update, _prepare_tp_inputs, and _prepare_tp_inputs_no_cache. Signed-off-by: ZhaoyangWang <zhaoyangw@nvidia.com>

zhaoyangwang-nvidia · 2026-06-02T09:01:23Z

/bot run

tensorrt-cicd · 2026-06-02T09:07:02Z

PR_Github #51578 [ run ] triggered by Bot. Commit: 9051c11 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>

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

zhaoyangwang-nvidia added 3 commits June 2, 2026 17:01

zhaoyangwang-nvidia force-pushed the merge-eagle-mtp branch from fde907a to 9051c11 Compare June 2, 2026 09:01

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

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!

zhaoyangwang-nvidia commented Jun 2, 2026

Uh oh!

tensorrt-cicd commented Jun 2, 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

zhaoyangwang-nvidia commented Mar 19, 2026 •

edited by coderabbitai Bot

Loading