[TRTLLM-14093][feat] Eagle3 support for MiniMax-M3 on the MSA backend

[zheyuf]

Overview

Adds one-model Eagle3 speculative decoding for MiniMax-M3 on the MSA
backend, with overlap scheduler and CUDA graph support.

This PR is for InferenceX Minimax-M3 bringup and is previous merged in side branch: brb-nv#4

What this enables

• 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 coalesces an index-K pool into its KV pages
  and exposes only synthetic AttentionOp tensors, which the dense Eagle3
  drafter cannot attend against — so it opts out of shared draft layers
  (supports_shared_draft_layers = False) and requires the separate draft
  manager even under attention DP.
• Draft KV manager defaults to V2 under V2 targets, with a draft page-size
  capability: the separate draft manager runs at tokens_per_block=32 to avoid
  a trtllm-gen IMA at 128 on the drafter's generation kernels, while the MSA
  target requires 128 (draft_manager_tokens_per_block).
• 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: 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.

Validation

All numbers below are from the source PR's runs (4×B200, TP4/EP4, Eagle3
draft_len=3, out-of-the-box config, no env vars). CI on this ported branch is
pending (/bot run).

Accuracy & acceptance — test_nvfp4_eagle3 (MSA + CUDA graphs + overlap scheduler)

Variant	MMLU	GSM8K	chat-GSM8K accept rate / length
attention_dp=False	85.14	90.94	0.843 / 3.528
attention_dp=True	85.11	89.95	0.819 / 3.456

Non-spec references: MMLU 85.11 / GSM8K 91.32 — accuracy-neutral.
Acceptance rate reference: 0.839 / 3.518 on vLLM https://huggingface.co/Inferact/MiniMax-M3-EAGLE3

Known issues (base-branch, not addressed here)

• Eagle3 drafter IMA at tokens_per_block=128: the drafter's trtllm-gen
  generation kernels hit an illegal memory access at page 128, independent of
  MSA. Worked around via the draft manager's page-size capability (draft runs at
  32); needs a kernel-team fix, after which the attribute can be removed.
• fmha_sm100 decode-planner cap: total_q × num_qo_heads ≤ 65536 bounds
  Eagle3 (4 verify tokens) at max_batch_size 1024 with TP4-sharded heads and 256
  under attention DP; the test caps accordingly.

Test plan

1. /bot run — builds the branch and runs the MSA+Eagle3 accuracy tests
   (test_nvfp4_eagle3, added to llm_function_core.txt).

[coderabbitai]

Walkthrough

This PR adds one-model Eagle3 speculative decoding support for MiniMax-M3. It updates the sparse KV cache manager's page-table and index handling, switches attention metadata to subclass TrtllmAttentionMetadata with revised extend/decode routing, wires spec_metadata through model layers, adds executor-side validation, adjusts dummy-request/draft-KV-cache handling, and adds accuracy/acceptance tests.

Changes

MiniMax-M3 Eagle3 Speculative Decoding

Layer / File(s)	Summary
Sparse cache manager overrides tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py	Overrides _prepare_page_table_tensor() for shared K/V/index pools with explicit pool mapping tensors, and reworks _get_batch_cache_indices_by_pool_id() to accept a block-count cap and directly slice base page indices instead of using base V1 conversion.
Attention metadata base class and routing tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py	Changes MiniMaxM3AttentionMetadata to subclass TrtllmAttentionMetadata, updates docs, and revises the extend-vs-decode predicate in prepare() to also trigger extend mode when any sequence adds more than one new token.
Decoder model spec_metadata wiring tensorrt_llm/_torch/models/modeling_minimaxm3.py	Adds optional spec_metadata parameters to MiniMaxM3DecoderLayer.forward and MiniMaxM3Model.forward for hidden-state capture, and switches MiniMaxM3ForCausalLM to subclass SpecDecOneEngineForCausalLM.
Executor validation and dummy-request handling tensorrt_llm/_torch/pyexecutor/py_executor_creator.py, tensorrt_llm/_torch/pyexecutor/model_engine.py, tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py, tensorrt_llm/runtime/kv_cache_manager_v2/_utils.py	Adds a NotImplementedError guard for unsupported minimax_m3 one-engine speculative configurations, skips attention warmup for minimax_m3, extends add_dummy_requests to support V1-family draft KV cache managers lacking _create_kv_cache, and adds a detailed assertion message to exact_div.
Eagle3 accuracy and acceptance tests tests/integration/defs/accuracy/references/gsm8k.yaml, tests/integration/defs/accuracy/references/mmlu.yaml, tests/integration/defs/accuracy/test_llm_api_pytorch.py	Adds MMLU/GSM8K accuracy reference entries for MiniMax-M3-NVFP4 with Eagle3 spec decoding, and a new integration test that runs accuracy evaluations plus a speculative acceptance-rate probe.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant MiniMaxM3Model
  participant MiniMaxM3DecoderLayer
  participant SpecMetadata
  participant TrtllmAttentionMetadata

  MiniMaxM3Model->>MiniMaxM3DecoderLayer: forward(hidden_states, residual, spec_metadata)
  MiniMaxM3DecoderLayer->>TrtllmAttentionMetadata: prepare() base KV lens/block offsets
  TrtllmAttentionMetadata-->>MiniMaxM3DecoderLayer: prepared metadata
  MiniMaxM3DecoderLayer->>SpecMetadata: maybe_capture_hidden_states(layer_idx, hidden_states, residual)
  MiniMaxM3DecoderLayer-->>MiniMaxM3Model: hidden_states

Loading

Possibly related PRs

• NVIDIA/TensorRT-LLM#12262: Threads SpecMetadata runtime fields that complement the new decoder-layer wiring introduced here.

Suggested reviewers: syuoni, mikeiovine, xxi-nv, xinhe-nv, yuxianq, jieli-matrix, yizhang-nv

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title matches the change set and follows the required [ticket][type] summary format.
Description check	✅ Passed	The description covers the overview, behavior changes, validation, known issues, and test plan required by the template.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py (1)

2684-2737: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Track V1 draft registrations in the cleanup path

The V1 draft-manager branch never sets draft_kv_cache, so a later kv_cache.resize() failure skips draft_kv_cache_manager.free_resources(current_request) even though the request was already registered via add_sequence_batch. That leaves the dummy request behind and can break the next warmup/retry with the same request id.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py` around lines 2684 -
2737, The V1 draft-manager path in the request setup/cleanup flow does not
register any handle for the draft allocation, so a later kv_cache.resize failure
bypasses freeing the already-added dummy request. Update the cleanup logic
around the draft_kv_cache_manager branch in the request preparation code to
track that a V1 draft registration happened and call
draft_kv_cache_manager.free_resources(current_request) on failure/retry paths,
alongside the existing draft_kv_cache handling.

🧹 Nitpick comments (1)

tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py (1)

2686-2686: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Prefer isinstance over hasattr duck-typing to distinguish manager type.

hasattr(draft_kv_cache_manager, "_create_kv_cache") distinguishes V2 vs V1-family managers via duck typing. Since this file defines KVCacheManagerV2 itself, isinstance(draft_kv_cache_manager, KVCacheManagerV2) would be a more robust and idiomatic check.

♻️ Proposed refactor

-                    if hasattr(draft_kv_cache_manager, "_create_kv_cache"):
+                    if isinstance(draft_kv_cache_manager, KVCacheManagerV2):

As per coding guidelines, "Prefer isinstance() to duck typing where possible; use built-in exception types for error handling, not custom exceptions in most cases."

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py` at line 2686, Replace
the duck-typing check in KVCacheManagerV2 with an explicit type check: use
isinstance(draft_kv_cache_manager, KVCacheManagerV2) instead of hasattr(...,
"_create_kv_cache"). Update the manager-selection branch in KVCacheManagerV2 to
rely on the concrete class name rather than probing for an internal method,
keeping the existing behavior for V2 vs V1-family managers intact.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py`:
- Around line 2684-2737: The V1 draft-manager path in the request setup/cleanup
flow does not register any handle for the draft allocation, so a later
kv_cache.resize failure bypasses freeing the already-added dummy request. Update
the cleanup logic around the draft_kv_cache_manager branch in the request
preparation code to track that a V1 draft registration happened and call
draft_kv_cache_manager.free_resources(current_request) on failure/retry paths,
alongside the existing draft_kv_cache handling.

---

Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py`:
- Line 2686: Replace the duck-typing check in KVCacheManagerV2 with an explicit
type check: use isinstance(draft_kv_cache_manager, KVCacheManagerV2) instead of
hasattr(..., "_create_kv_cache"). Update the manager-selection branch in
KVCacheManagerV2 to rely on the concrete class name rather than probing for an
internal method, keeping the existing behavior for V2 vs V1-family managers
intact.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: e8a2f3cc-ff2e-4d5b-bcf3-c6b07c3fb51e

📥 Commits

Reviewing files that changed from the base of the PR and between a0c406f and c636333.

📒 Files selected for processing (10)

• tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py
• tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/metadata.py
• tensorrt_llm/_torch/models/modeling_minimaxm3.py
• tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
• tensorrt_llm/_torch/pyexecutor/model_engine.py
• tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
• tensorrt_llm/runtime/kv_cache_manager_v2/_utils.py
• tests/integration/defs/accuracy/references/gsm8k.yaml
• tests/integration/defs/accuracy/references/mmlu.yaml
• tests/integration/defs/accuracy/test_llm_api_pytorch.py

[zheyuf]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61126 [ run ] triggered by Bot. Commit: d245591 Link to invocation

[tensorrt-cicd]

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

[zheyuf]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61401 [ run ] triggered by Bot. Commit: b077bc4 Link to invocation

[tensorrt-cicd]

PR_Github #61401 [ run ] completed with state SUCCESS. Commit: b077bc4
/LLM/main/L0_MergeRequest_PR pipeline #49628 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

[zheyuf]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61439 [ run ] triggered by Bot. Commit: b077bc4 Link to invocation

[tensorrt-cicd]

PR_Github #61439 [ run ] completed with state SUCCESS. Commit: b077bc4
/LLM/main/L0_MergeRequest_PR pipeline #49663 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

[nvpohanh]

[by Codex] @liji-nv Could you please review this PR? Thank you!