[TRTLLM-14093][feat] One-model Eagle3 speculative decoding for MiniMax-M3 on the MSA backend - #4
Merged
Merged
Conversation
zheyuf
force-pushed
the
feat/m3_with_msa-eagle3
branch
from
July 18, 2026 05:55
06204aa to
13642d2
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> (cherry picked from commit 448c489) Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
…V2 targets A structurally-V2 target (sparse-attention managers like MiniMax-M3's, which don't set use_kv_cache_manager_v2) paired with a plain-transformer draft (Eagle3) resolved the draft manager to V1, requiring V1-family special cases in KVCacheManagerV2.add_dummy_requests and leaving a latent AttributeError: KVCacheV2Scheduler calls suspend_request() on the draft manager, which only exists on V2. Promote the draft class to KVCacheManagerV2 whenever the target is V2 (shared helper used by both the creation and the cache-cost estimation paths, which previously disagreed). add_dummy_requests reverts to its original V2-only shape - kv_cache_manager_v2.py returns byte-identical to its pre-enablement upstream state. The V1-draft branches were reachable by exactly one configuration - MiniMax-M3 + one-model Eagle3; every other V2-target spec config already resolves a V2 draft (flag or target-config fallback) or shares the target manager. The V2-draft combination was validated on 4xB200: MMLU 84.84 / GSM8K 90.07 (explicit flag), plus boot+acceptance probe on the new default path (AR 0.465/AL 2.40, matching the V1-draft band). Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com> (cherry picked from commit 3854f4d) Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
… verify Removes the MSA+speculative-decoding restriction: the MSA kernel path (sparse_use_msa=True) now verifies draft tokens through a multi-token decode driver instead of rejecting at creation. Driver (decode_wrapper): hybrid scheme, validated bit-exact against the eager fmha_sm100 api. The proxy (OnlyScore) pass runs natively multi-token: the kernel's inclusive causal bound with offset = kv_len - qo_len produces exactly the verify ladder (token t attends kv_len - qo_len + t + 1 positions). Top-k selection gets per-token ladder valid-page counts, so a draft token cannot select blocks past its own attend bound. The sparse GQA pass row-expands each token to a qo_len=1 pseudo-row keeping its request's full kv_len and page-table base, with the ladder in the per-row offset - the same transform the eager api applies internally; the kernel consumes kv_block_indexes per row, so native multi-token sparse is not expressible. All device ops, capture-safe. Routing/metadata: on MSA, pure-generation uniform multi-token batches stay DECODE-shaped (decode_qo_len = 1 + draft_len), keeping the captured/overlap-safe device-plan path; mixed context+gen batches take the eager extend path; the reference backend is unchanged. KV slot staging generalizes to the causal ladder, and dense layers 0-2 get a ladder SDPA mask (identical to the old math at qo_len=1). Overlap-correction fix: the MSA hook now re-stages the flat page table on mixed batches, not just the CPU length mirrors - a correction that shrinks a row across a page boundary changes the indptr layout the eager kernels rebuild from the corrected lens, misbasing every subsequent row's pages (symptom: MMLU passes, GSM8K collapses). Tests: driver bit-diff suite extended with qo_len=4 legs - proxy bit-diff vs the eager api, top-k vs a per-token ladder reference, and CUDA-graph capture/replay with mutated lens. test_nvfp4_eagle3 runs the MSA path (use_msa single-choice parametrize); QA-list rows updated to the new test IDs (the old rows no longer matched any collected test). Validation (4xB200, TP4/EP4, NVFP4): MSA adp=False MMLU 84.97 / GSM8K 90.18, AR 0.698 / AL 3.094; MSA adp=True 85.21 / 91.24, AR 0.718 / AL 3.153; reference regression control 85.04 / 90.33, AR 0.707 / AL 3.121. MSA suite wall-clock 6:10 vs reference 11:50 on the same gates. Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com> (cherry picked from commit 4d44c58) Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
The MSA decode driver was built capture-safe (device-only replan, prepare-time buffer allocation, stable data_ptrs), so multi-token verify no longer routes through the eager extend path on the MSA backend. Allow cuda_graph_config with sparse_use_msa=True; the reference path keeps raising. - speculative/eagle3.py: save/restore kv_lens_cuda across CUDA-graph warmup iterations. The one-model worker saved _seq_lens and the spec-decoding tensors, but not kv_lens_cuda, which the draft loop mutates in place - so the runner's pre-capture warmup iterations ran with drifted kv lens. Same pattern dflash/pard use; active only during graph warmup, not capture. - py_executor_creator.py: the graphs+spec rejection now applies only to the reference path (its verify routes through the eager extend path). - Capture hardening: the dense layers 0-2 decode branch now expands GQA K/V per KV head instead of all heads at once (bitwise-identical math; the expansion is captured into the graph pool, and under attention DP - unsharded heads - the whole-tensor transient exceeds the pool budget at large graph buckets). They also baked the prepare-time host upper bound max_seqlen_k into the captured gather/mask width; replays whose kv_len outgrew the capture-time value would silently truncate attention. Under graphs, bake min(page-table capacity, engine max_seq_len) instead (raw capacity alone OOMs during KV estimation; the seq_lens mask already invalidates positions past each row's true length). resolve_decode_state now raises if it would BUILD a decode state (JIT + allocation) while the stream is capturing. - test_nvfp4_eagle3: cuda_graph single-choice parametrize; the gated variants run the endgame config (MSA + Eagle3 + overlap + graphs). Support matrix updated (EAGLE-3 Linear: Yes). Probe (4xB200, TP4/EP4, batch-1, overlap, greedy): 25.98 tok/s with graphs vs 13.70 eager (1.90x), AR/AL in band; accuracy gate under graphs (adp=False): MMLU 84.62 / GSM8K 90.86, AR 0.719 / AL 3.156 - identical to eager; spec-dec graph capture confirmed on all ranks (draft_len=3 buckets); negative control: reference path + graphs + spec still raises at creation. Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com> (cherry picked from commit c8b09eb) Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
zheyuf
force-pushed
the
feat/m3_with_msa-eagle3
branch
from
July 20, 2026 06:27
2af9ecc to
b52ada8
Compare
Squash of the follow-up work that makes the four ported Eagle3 commits functional and validated on feat/m3_with_msa: - Adapt to this branch's API: the creator guard and integration test used the removed sparse_use_msa field (implementation= now); MsaSparseGqaFmha.is_available crashed on the dense Eagle3 drafter (sparse_params=None). - Correct the MSA length sources: msa_kv_lens_cpu excludes num_extra_kv_tokens (nonzero only under speculation; inflated every slot/page/plan length and crashed spec warmup with an illegal memory access), and the graph-safe plan mirrors are sized per expanded token row (the planner splits qo_len > 1 requests per token). - Device-side overlap correction: on_update_kv_lens on the MSA metadata re-derives the KV-write slots, per-token valid-block counts, and the per-row kv_segment_lens/qo_offset plan-mirror entries from the corrected kv_lens_cuda — pure in-stream device ops, capture-safe, shrink-only. Eager mixed batches install corrected host lens and rebuild fields and eager plans. Without this the overlap scheduler (on by default) silently corrupts the KV layout on every partially rejected verify step (GSM8K 57.9 vs 90.4+ with the hook). - Draft KV page size: MiniMaxM3KVCacheManagerV2 declares draft_manager_tokens_per_block = 32, read generically at draft-manager creation — a scoped workaround for a base regression (the drafter's trtllm-gen generation kernels IMA at tokens_per_block=128; the MSA target requires 128, the dense draft layers have no page-size constraint). Remove once the kernel bug is fixed. - Tests: test_nvfp4_eagle3 runs the endgame config (MSA + Eagle3 + overlap + CUDA graphs, attention-DP on/off) with MMLU/GSM8K gates and a chat-format GSM8K acceptance check against the drafter head's published reference (floors 0.78 / 3.3; head card 0.839 / 3.518); QA-list rows updated; two unit tests for token-count scratch sizing and the per-token valid-block ladder. - Review feedback: creator-level M3 spec guards removed, ADP shared- layout opt-out documented on the manager class, complexity/comment pass (net -70 lines), pre-commit formatting. Validated on 4xB200 (TP4/EP4, NVFP4, draft_len=3, no env setup): adp=False MMLU 85.14 / GSM8K 90.94, chat acceptance 0.843 / 3.528; adp=True 84.89 / 90.37, chat acceptance 0.846 / 3.538; spec-off regression test_nvfp4 both variants pass; 42/42 unit tests; overlap throughput +6.2-6.4% at acceptance-matched batches 32/128. Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>
zheyuf
force-pushed
the
feat/m3_with_msa-eagle3
branch
from
July 20, 2026 06:41
b52ada8 to
499a369
Compare
zheyuf
marked this pull request as ready for review
July 20, 2026 07:20
zheyuf
added a commit
to zheyuf/TensorRT-LLM
that referenced
this pull request
Jul 22, 2026
…x-M3 on the MSA backend Port of brb-nv#4 (head 499a369) onto main. - Enable one-model Eagle3 for MiniMax-M3: spec-metadata capture hooks in the decoder layers, SpecDecOneEngineForCausalLM base, and a separate draft KV cache manager (the MSA manager's INDEX_KEY-coalesced pools expose synthetic AttentionOp tensors dense draft layers cannot share, so it opts out of shared draft layers even under attention DP). - Default the one-model draft KV manager to V2 under V2 targets, with a draft page-size capability (the draft manager runs at tokens_per_block=32 to avoid a trtllm-gen IMA at 128 on the drafter's generation kernels; the MSA target requires 128). - MSA multi-token decode for spec verify: per-token plan rows, causal ladder masks in the dense SDPA path, and proxy scratch sized by the worst-case decode token count instead of the batch size. - CUDA graphs + overlap scheduler support: graph-safe max_k baking, an on_update_kv_lens hook that re-derives MSA slots/bounds on device from the corrected kv_lens, and kv_lens_cuda save/restore during graph warmup. Conflicts against main's revised MSA backend (NVIDIA#16291) were resolved by keeping main's _build_decode_plans naming and per-layer eager planning; the side branch's eager-plan dedup and HND index-K tests never landed on main and are not resurrected here. Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>
zheyuf
added a commit
to zheyuf/TensorRT-LLM
that referenced
this pull request
Jul 23, 2026
…x-M3 on the MSA backend Port of brb-nv#4 (head 499a369) onto main. - Enable one-model Eagle3 for MiniMax-M3: spec-metadata capture hooks in the decoder layers, SpecDecOneEngineForCausalLM base, and a separate draft KV cache manager (the MSA manager's INDEX_KEY-coalesced pools expose synthetic AttentionOp tensors dense draft layers cannot share, so it opts out of shared draft layers even under attention DP). - Default the one-model draft KV manager to V2 under V2 targets, with a draft page-size capability (the draft manager runs at tokens_per_block=32 to avoid a trtllm-gen IMA at 128 on the drafter's generation kernels; the MSA target requires 128). - MSA multi-token decode for spec verify: per-token plan rows, causal ladder masks in the dense SDPA path, and proxy scratch sized by the worst-case decode token count instead of the batch size. - CUDA graphs + overlap scheduler support: graph-safe max_k baking, an on_update_kv_lens hook that re-derives MSA slots/bounds on device from the corrected kv_lens, and kv_lens_cuda save/restore during graph warmup. Conflicts against main's revised MSA backend (NVIDIA#16291) were resolved by keeping main's _build_decode_plans naming and per-layer eager planning; the side branch's eager-plan dedup and HND index-K tests never landed on main and are not resurrected here. Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>
brb-nv
pushed a commit
that referenced
this pull request
Jul 23, 2026
…x-M3 on the MSA backend (#4) Signed-off-by: Zheyu Fu <zheyuf@nvidia.com> Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
zheyuf
added a commit
to zheyuf/TensorRT-LLM
that referenced
this pull request
Jul 23, 2026
…x-M3 on the MSA backend Port of brb-nv#4 (head 499a369) onto main. - Enable one-model Eagle3 for MiniMax-M3: spec-metadata capture hooks in the decoder layers, SpecDecOneEngineForCausalLM base, and a separate draft KV cache manager (the MSA manager's INDEX_KEY-coalesced pools expose synthetic AttentionOp tensors dense draft layers cannot share, so it opts out of shared draft layers even under attention DP). - Default the one-model draft KV manager to V2 under V2 targets, with a draft page-size capability (the draft manager runs at tokens_per_block=32 to avoid a trtllm-gen IMA at 128 on the drafter's generation kernels; the MSA target requires 128). - MSA multi-token decode for spec verify: per-token plan rows, causal ladder masks in the dense SDPA path, and proxy scratch sized by the worst-case decode token count instead of the batch size. - CUDA graphs + overlap scheduler support: graph-safe max_k baking, an on_update_kv_lens hook that re-derives MSA slots/bounds on device from the corrected kv_lens, and kv_lens_cuda save/restore during graph warmup. Conflicts against main's revised MSA backend (NVIDIA#16291) were resolved by keeping main's _build_decode_plans naming and per-layer eager planning; the side branch's eager-plan dedup and HND index-K tests never landed on main and are not resurrected here. Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>
zheyuf
added a commit
to zheyuf/TensorRT-LLM
that referenced
this pull request
Jul 24, 2026
…x-M3 on the MSA backend Port of brb-nv#4 (head 499a369) onto main. - Enable one-model Eagle3 for MiniMax-M3: spec-metadata capture hooks in the decoder layers, SpecDecOneEngineForCausalLM base, and a separate draft KV cache manager (the MSA manager's INDEX_KEY-coalesced pools expose synthetic AttentionOp tensors dense draft layers cannot share, so it opts out of shared draft layers even under attention DP). - Default the one-model draft KV manager to V2 under V2 targets, with a draft page-size capability (the draft manager runs at tokens_per_block=32 to avoid a trtllm-gen IMA at 128 on the drafter's generation kernels; the MSA target requires 128). - MSA multi-token decode for spec verify: per-token plan rows, causal ladder masks in the dense SDPA path, and proxy scratch sized by the worst-case decode token count instead of the batch size. - CUDA graphs + overlap scheduler support: graph-safe max_k baking, an on_update_kv_lens hook that re-derives MSA slots/bounds on device from the corrected kv_lens, and kv_lens_cuda save/restore during graph warmup. Conflicts against main's revised MSA backend (NVIDIA#16291) were resolved by keeping main's _build_decode_plans naming and per-layer eager planning; the side branch's eager-plan dedup and HND index-K tests never landed on main and are not resurrected here. Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>
zheyuf
added a commit
to zheyuf/TensorRT-LLM
that referenced
this pull request
Jul 26, 2026
…x-M3 on the MSA backend Port of brb-nv#4 (head 499a369) onto main. - Enable one-model Eagle3 for MiniMax-M3: spec-metadata capture hooks in the decoder layers, SpecDecOneEngineForCausalLM base, and a separate draft KV cache manager (the MSA manager's INDEX_KEY-coalesced pools expose synthetic AttentionOp tensors dense draft layers cannot share, so it opts out of shared draft layers even under attention DP). - Default the one-model draft KV manager to V2 under V2 targets, with a draft page-size capability (the draft manager runs at tokens_per_block=32 to avoid a trtllm-gen IMA at 128 on the drafter's generation kernels; the MSA target requires 128). - MSA multi-token decode for spec verify: per-token plan rows, causal ladder masks in the dense SDPA path, and proxy scratch sized by the worst-case decode token count instead of the batch size. - CUDA graphs + overlap scheduler support: graph-safe max_k baking, an on_update_kv_lens hook that re-derives MSA slots/bounds on device from the corrected kv_lens, and kv_lens_cuda save/restore during graph warmup. Conflicts against main's revised MSA backend (NVIDIA#16291) were resolved by keeping main's _build_decode_plans naming and per-layer eager planning; the side branch's eager-plan dedup and HND index-K tests never landed on main and are not resurrected here. Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>
zheyuf
added a commit
to zheyuf/TensorRT-LLM
that referenced
this pull request
Jul 27, 2026
…x-M3 on the MSA backend Port of brb-nv#4 (head 499a369) onto main. - Enable one-model Eagle3 for MiniMax-M3: spec-metadata capture hooks in the decoder layers, SpecDecOneEngineForCausalLM base, and a separate draft KV cache manager (the MSA manager's INDEX_KEY-coalesced pools expose synthetic AttentionOp tensors dense draft layers cannot share, so it opts out of shared draft layers even under attention DP). - Default the one-model draft KV manager to V2 under V2 targets, with a draft page-size capability (the draft manager runs at tokens_per_block=32 to avoid a trtllm-gen IMA at 128 on the drafter's generation kernels; the MSA target requires 128). - MSA multi-token decode for spec verify: per-token plan rows, causal ladder masks in the dense SDPA path, and proxy scratch sized by the worst-case decode token count instead of the batch size. - CUDA graphs + overlap scheduler support: graph-safe max_k baking, an on_update_kv_lens hook that re-derives MSA slots/bounds on device from the corrected kv_lens, and kv_lens_cuda save/restore during graph warmup. Conflicts against main's revised MSA backend (NVIDIA#16291) were resolved by keeping main's _build_decode_plans naming and per-layer eager planning; the side branch's eager-plan dedup and HND index-K tests never landed on main and are not resurrected here. Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>
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.
Overview
Support Eagle3 + overlap scheduler + cuda graph on MSA backends.
Commits
[feat] Enable one-model Eagle3 speculative decoding for MiniMax-M3
[chore] Default one-model draft KV manager to V2 under V2 targets
[feat] MSA multi-token decode for one-model Eagle3 spec verify
[feat] Enable CUDA graphs for MiniMax-M3 MSA + Eagle3
[feat] Adapt Eagle3 support to the new side branch
Validation
All runs: 4xB200, TP4/EP4, NVFP4 target, Eagle3 draft_len=3, out-of-the-box
config (no env vars required).
Accuracy & acceptance —
test_nvfp4_eagle3(MSA + CUDA graphs + overlap scheduler)All rows at this PR's head (base
ec6c531752):Non-spec references: MMLU 85.11 / GSM8K 91.32 — accuracy-neutral.
Acceptance parity with the drafter head's published reference
Chat-format GSM8K n=200 (the head card's own setup,
https://huggingface.co/Inferact/MiniMax-M3-EAGLE3 → 0.839 / 3.518 on vLLM):
this PR measures 0.82–0.85 / 3.46–3.54 across six runs spanning MXFP8/NVFP4
targets, MSA/triton backends, and attention-DP on/off — parity with the
reference. The test gates on floors 0.78 / 3.3.
Known issues (base-branch, not addressed by this PR)
tokens_per_block=128: the drafter's trtllm-gengeneration kernels hit an illegal memory access at page 128 — independent
of MSA (the triton path + Eagle3 forced to page 128 crashes identically;
clean at 32; the same config was green on the pre-rework base). Worked
around via the draft manager page-size capability above; needs a
kernel-team fix, after which the capability attribute can be removed.
total_q x num_qo_heads <= 65536boundsEagle3 (4 verify tokens) at max_batch_size 1024 with TP4-sharded heads and
256 under attention DP; the test caps accordingly.
Test plan
/bot run— builds the branch, runs the MSA+Eagle3 accuracy tests inllm_function_core.txt(no environment setup needed).