[TRTLLM-14093][feat] Eagle3 for M3 non-MSA path + overlap scheduler compatibility - #2
Merged
Conversation
zheyuf
force-pushed
the
feat/m3-eagle3-brb-port
branch
from
July 9, 2026 06:40
90be465 to
aab4935
Compare
… MiniMax-M3 Squashed port of feat/minimax-m3-eagle3 (PR NVIDIA#16021) onto feat/branch_m3. Target-side enablement for the Inferact/MiniMax-M3-EAGLE3 draft head on the reference (non-MSA) sparse backend: - Wire MiniMaxM3ForCausalLM as SpecDecOneEngineForCausalLM and capture Eagle3 aux hidden states at layer exit (fully TP-reduced; no cross-layer allreduce+norm fusion). - Rebase MiniMaxM3AttentionMetadata onto TrtllmAttentionMetadata (the shared per-step metadata is consumed by TRTLLM draft layers, the Eagle3 one-model worker, and engine isinstance gates; precedent: DSA and the M3 MSA metadata). - Route multi-token generation rows (spec verify: 1 + draft_len tokens) through the extend path; the decode branch stays reserved for batches where every row appends exactly one token. - Vectorized sync-free builder shared with a new on_update_kv_lens hook that re-derives seq_lens/prefix_lens/q_positions/out_cache_loc on device from the corrected kv_lens_cuda under overlap scheduler + spec (DSA pattern); in-bounds clamps cover the optimistic page-boundary overhang (max_seqlen_k SDPA width + slot gathers). - V1-family draft manager support in KVCacheManagerV2.add_dummy_requests (exception-safe); attention-DP dummy requests register in the draft manager; MiniMaxM3KVCacheManagerV2 opts out of shared draft layers under attention DP (its AttentionOp tensors are synthetic). - Creation-time guards: tree modes, disabled separate draft KV (disagg WAR), CUDA graphs, and MSA+spec (the in-builder MSA rejection is hoisted above routing so mixed batches cannot bypass it). - Accuracy test test_nvfp4_eagle3 (MMLU + GSM8K + acceptance probe, attention_dp parametrized) + reference rows. Validated on 4xB200 (NVFP4 tp4/ep4, draft_len=3, overlap scheduler on, eager) at this commit: - test_nvfp4_eagle3[attention_dp=False]: MMLU 85.50 / GSM8K 89.73 (refs 83/88), acceptance rate 0.709 / mean acceptance length 3.126 - test_nvfp4_eagle3[attention_dp=True]: MMLU 85.14 / GSM8K 91.32, acceptance rate 0.767 / mean acceptance length 3.302 - batch-1 greedy: 6.05 -> 17.05 tok/s (2.82x) - spec-off boot + generation clean (TRTLLM-Gen warmup now runs for M3 and is harmless) Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>
zheyuf
changed the title
zheyuf
force-pushed
the
feat/m3-eagle3-brb-port
branch
from
July 9, 2026 07:46
aab4935 to
448c489
Compare
zheyuf
marked this pull request as ready for review
July 9, 2026 07:52
zheyuf
changed the title
zheyuf
changed the title
zheyuf
changed the title
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
1 participant
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.
Squashed, re-validated port of NVIDIA#16021 onto
feat/branch_m3: one-model Eagle3 speculative decoding for MiniMax-M3 on the reference (non-MSA) sparse backend, pairing with theInferact/MiniMax-M3-EAGLE3draft head.Measured on this commit (4xB200, NVFP4 tp4/ep4,
max_draft_len=3, overlap scheduler on, eager):test_nvfp4_eagle3[attention_dp=False-overlap_scheduler=True]test_nvfp4_eagle3[attention_dp=True-overlap_scheduler=True](Reference thresholds 83 / 88; non-spec baselines 85.11 / 91.32 — accuracy-neutral.)
Batch-1 greedy throughput: 6.05 -> 17.05 tok/s (2.82x).
Limitations
sparse_use_msa=Trueis rejected at creation time; Eagle3 currently runs on the reference (non-MSA) sparse backend. Multi-token support in the MSA decode driver is the follow-up that would lift this (the underlying kernel's pack factors already allowqo_len <= 32).use_kv_cache_manager_v2=Trueis mechanically supported by the dummy-request dispatch but not yet validated for MiniMax-M3. V1 vs V2 for the drafter is not expected to make a performance difference.cuda_graph_configmust be null with Eagle3; all numbers above are eager): multi-token verify routes through the M3 extend path, which is not capture-safe yet (host-side reads in the dense-layer fallback). Non-speculative M3 keeps its existing CUDA-graph support. Likely follow-up paths: piecewise capture, or the multi-token MSA decode driver (which is already capture-safe by design).