[TRTLLM-14093][feat] Eagle3 support for MiniMax-M3 on the MSA backend (consolidated, staging for #16021)

[zheyuf]

Staging PR for self-review. This is the consolidated MiniMax-M3 Eagle3 port rebased on main (NVIDIA/TensorRT-LLM@a67ede16f, 2026-09-07). Once it is reviewed here, this branch will be force-pushed to feat/minimax-m3-eagle3 (NVIDIA#16021); no CI runs on this fork, /bot run happens there. The two commits are meant to be reviewed in order: (1) Eagle3 on the MSA backend, (2) the shared draft KV cache (virtual attention-op pools) + tests.

Overview

One-model Eagle3 speculative decoding for MiniMax-M3 on the MSA sparse-attention backend, with CUDA graphs, the overlap scheduler, aggregated and disaggregated serving. This consolidates the Eagle3 work validated on the feat/m3_with_msa side branch — NVIDIA#16021 (this PR's original content), NVIDIA#17341, NVIDIA#17457, NVIDIA#17656 and NVIDIA#18066 — into the design those follow-ups converged on, rebased on current main.

What this enables

• Eagle3 on the MSA backend. Spec-metadata capture hooks in the decoder layers, SpecDecOneEngineForCausalLM as the model base, multi-token (1 + draft_len) decode/verify in the MSA metadata: per-token cache slots, valid-block counts and plan rows, proxy scratch sized by the worst-case decode token count.
• Overlap scheduler + CUDA graphs. on_update_kv_lens re-derives cache slots, valid-block counts and the fmha_sm100 plans' per-row length mirrors on device after the scheduler corrects kv_lens, clamped to the staged lengths: sync-free, capture-safe, idempotent. The clamp also covers the draft loop advancing kv_lens_cuda between CUDA-graph warmup runs. The staging and the patch only run when speculative decoding is active, so non-speculative MiniMax-M3 steps are unchanged.
• Unified draft KV cache. MiniMax-M3 keeps its Eagle3 draft layer in the target KV cache manager in every supported configuration, so block reuse, eviction and disaggregated KV transfer cover the drafter's KV natively. The draft layer runs the generic TRTLLM attention op, which cannot address M3's non-uniform mega-slot through the regular pool mapping (KVCacheManagerV2 groups sub-pages by size, and at M3's production geometry the index-K page coalesces into the K/V pool). The M3 manager therefore presents each shared draft layer to the attention op as its own virtual pool rooted at the layer's K page, the same pattern DeepseekV4CacheManager and the SWA scratch-reuse path use. The draft loop runs on the shared manager like every other unified-KV model: no draft-side manager, no metadata swap, no change to the speculative-decoding code. trtllm-gen accepts the P128 draft shapes via an opt-in consulted only after its allowlist rejects the page size.
• Dense reference path. The Triton/SDPA metadata gains the causal-ladder verify mask and device-side slot derivation.

Differences from the side-branch commits

• The separate draft manager path (supports_shared_draft_layers=False, draft_manager_tokens_per_block=32, the draft-manager V2 promotion helper, the dummy-request draft-manager pass-through) is not carried: [None][perf] wrap up eagle changes in m3 side branch NVIDIA/TensorRT-LLM#18066 retired it, and main already promotes draft managers via _validate_or_fallback_kv_cache_manager_v2.
• The P32 sub-page view ([None][feat] MiniMax-M3 eagle+disagg pipe clean: share Eagle3 draft layers via a sub-page view NVIDIA/TensorRT-LLM#17457/[None][fix] Correct MiniMax-M3 Eagle draft sub-page pool bound NVIDIA/TensorRT-LLM#17656) and the native P128 draft view ([None][perf] wrap up eagle changes in m3 side branch NVIDIA/TensorRT-LLM#18066) are replaced by virtual attention-op pools inside the M3 manager; the draft-manager resolution hooks the side branch added to speculative/utils.py and SpecWorkerBase are not needed.
• The Eagle hidden-state publication path from [None][perf] wrap up eagle changes in m3 side branch NVIDIA/TensorRT-LLM#18066 (eagle_hidden_states_copy, multi-stream DAG edge) is not ported: it was shown to be inert on the piecewise context path.
• Hybrid NVFP4 KV cache pieces ([None][perf] Add MiniMax-M3 hybrid NVFP4 KV cache support NVIDIA/TensorRT-LLM#17985) are out of scope for this PR.
• Review feedback on the original diff: _msa_page_size, msa_qo_lens_dev and the live batch/token counters are gone (the metadata's tokens_per_block, seq_lens_cuda, num_seqs/num_tokens are used); mixed steps no longer rebuild with a D2H copy.
• Shared-path footprint: the trtllm-gen P128 opt-in (flashinfer_trtllm_gen.py, consulted only after the allowlist rejects the page size) and the one-line MiniMax-M3 policy in py_executor_creator.py. speculative/, kv_cache_manager_v2.py, _util.py, eagle3.py and the other executor files are untouched; for every model other than MiniMax-M3 the two touched paths return exactly what main returns today.
• Needs [https://nvbugs/6698606][fix] Keep quantized KV context on TRTLLM-Gen FMHA NVIDIA/TensorRT-LLM#18545 (merged 2026-09-03, included in this base): it keeps FP8/NVFP4-KV context attention on trtllm-gen. Without it main routes the draft layer's BF16 context (small batches everywhere, all batches on SM103) to the C++ attention op; on a wheel built for SM100 only that op finds no SM103 kernel and falls back to unfused MHA, whose O(seq_len^2) workspace OOMs on MiniMax-M3 (see [None][fix] build SM103 cubins in the local perf harness wheel NVIDIA/TensorRT-LLM#18831 for the SM103-cubin gap in local wheels).
• Why the virtual pools exist and when they retire is written on MiniMaxM3KVCacheManagerV2: once V2 can keep a buffer role in its own pool, the K/V pool is uniform again and the virtual pools plus the P128 opt-in can be deleted.

Test Coverage

• TestMiniMaxM3::test_nvfp4_eagle3 (GQA Eagle3 head, MMLU + GSM8K + chat-GSM8K acceptance probe), aggregated arms with/without attention DP and a disaggregated (context TP2 -> generation TP2, NIXL) arm; scheduled in l0_dgx_b200 and the QA list.
• Unit tests (CPU): multi-token valid-block counts and proxy scratch sizing, on_update_kv_lens slot/count/plan-mirror re-derivation (incl. mixed plans, the staged-length clamp and the speculative-decoding gate), the virtual attention-op pool tables and the block-offset copy into them, shared-draft layout derivation.

Validation on main (this branch, GB300 TP4/EP4, MSA, FP8 KV, CUDA graphs, overlap scheduler)

• Unit tests: 36 passed (tests/unittest/_torch/attention/sparse/test_minimax_m3_*.py).
• Eagle3 draft_len=3, GQA head, 200 chat-GSM8K prompts, greedy: acceptance rate 0.832, mean acceptance length 3.50 (test floors 0.80 / 3.4; drafter card 0.839 / 3.518). GSM8K exact match 0.860 with Eagle3 vs 0.850 for the plain target on the same prompts.

Validation (side branch, same design)

4xB200, TP4/EP4, Eagle3 draft_len=3, GQA head: chat-GSM8K acceptance 0.832-0.835 / mean acceptance length 3.50 aggregated (ADP off/on), 0.827-0.828 / 3.48 disaggregated (NVIDIA#18066). Reference: 0.839 / 3.518 on vLLM (Inferact/MiniMax-M3-EAGLE3 card).

PR Checklist

• PR title follows the [JIRA/NVBUG/None][type] format
• Commits are signed off (DCO)
• ruff (new-style files) / yapf + isort (legacy files) clean
• CI (/bot run)