From e9d76f93e1252878984069b1a55c779ef3a50ac6 Mon Sep 17 00:00:00 2001
From: Zheyu Fu <zheyuf@NVIDIA.com>
Date: Wed, 22 Jul 2026 22:56:36 +0000
Subject: [PATCH 1/2] [TRTLLM-14093][feat] One-model Eagle3 speculative
 decoding for MiniMax-M3 on the MSA backend

Port of brb-nv/TensorRT-LLM#4 (head 499a3699ee) 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 (#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>
---
 docs/source/models/supported-models.md        |   4 +-
 .../sparse/minimax_m3/cache_manager.py        |  14 ++
 .../sparse/minimax_m3/msa_backend.py          | 203 +++++++++++++++---
 .../sparse/minimax_m3/triton_metadata.py      | 168 ++++++++++-----
 .../_torch/models/modeling_minimaxm3.py       | 104 +++++----
 tensorrt_llm/_torch/pyexecutor/_util.py       |  49 ++++-
 tensorrt_llm/_torch/pyexecutor/py_executor.py |   7 +
 tensorrt_llm/_torch/speculative/eagle3.py     |  15 +-
 .../defs/accuracy/references/gsm8k.yaml       |   3 +
 .../defs/accuracy/references/mmlu.yaml        |   3 +
 .../defs/accuracy/test_llm_api_pytorch.py     |  99 +++++++++
 .../test_lists/qa/llm_function_core.txt       |   2 +
 .../sparse/test_minimax_m3_msa_backend.py     |  24 +++
 13 files changed, 563 insertions(+), 132 deletions(-)

diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md
index b03f9961f7cc..38849627d08a 100644
--- a/docs/source/models/supported-models.md
+++ b/docs/source/models/supported-models.md
@@ -78,7 +78,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl
 | `Gemma4ForConditionalGeneration` | Untested          | Yes        | Untested                   | No                    | Yes             | No  | No               | No                | No     | Yes           | Untested         | No             | Yes                      | Untested              | Untested        |
 | `Gemma4UnifiedForConditionalGeneration` | Untested          | Untested   | Untested                   | No                    | Yes             | No  | No               | No                | No     | Yes           | Untested         | No             | Yes                      | Untested              | Untested        |
 | `Step3p7ForConditionalGeneration`| Yes               | Yes        | Yes                        | Untested              | Untested        | Yes | No               | No                | No     | Yes           | Untested         | Untested       | Yes                      | Untested              | Untested        |
-| `MiniMaxM3SparseForConditionalGeneration` [^12] | Yes               | Yes        | Yes                        | Untested              | Untested        | No  | No               | No                | No     | Yes           | Untested         | No             | N/A                      | Untested              | Untested        |
+| `MiniMaxM3SparseForConditionalGeneration` [^12] | Yes               | Yes        | Yes                        | Untested              | Untested        | No  | Yes              | No                | No     | Yes           | Untested         | No             | N/A                      | Untested              | Untested        |
 
 [^1]: Chunked Prefill for MLA can only be enabled on SM90/SM100/SM103/SM120.
 [^2]: KV cache reuse for MLA can only be enabled on SM90/SM100/SM103/SM120/SM121 and in BF16/FP8 KV cache dtype.
@@ -90,7 +90,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl
 [^9]: Audio modality only supported on E2B/E4B variants.
 [^10]: Audio requires a checkpoint with a `sound_config` and is supported only on the full (non-disaggregated) model path, not the EPD disaggregated path.
 [^11]: DeepSeek-V4 is only supported on Blackwell GPUs (`SM100+`). See the [DeepSeek-V4 example README](../../../examples/models/core/deepseek_v4/README.md) for setup and parallelism.
-[^12]: Supports text, image, and video inputs over the block-sparse attention path. The published MXFP8 checkpoint is dequantized on load so the runtime sees an effectively BF16 model. The text decoder is also usable standalone (text-only) via the `MiniMaxM3SparseForCausalLM` architecture. KV cache reuse and MTP are not supported on the sparse-attention path in this release.
+[^12]: Supports text, image, and video inputs over the block-sparse attention path. The published MXFP8 checkpoint is dequantized on load so the runtime sees an effectively BF16 model. The text decoder is also usable standalone (text-only) via the `MiniMaxM3SparseForCausalLM` architecture. KV cache reuse and MTP are not supported on the sparse-attention path in this release. One-model linear EAGLE-3 is supported; combining it with CUDA graphs requires the MSA kernels (`sparse_use_msa=True`, SM100).
 [^13]: The Cosmos 3 family also supports visual generation through the VisualGen API. See [Visual Generation Models](#visual-generation-models).
 
 # Multimodal Feature Support Matrix (PyTorch Backend)
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py
index 65c89b170c5c..26c20b114c93 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/cache_manager.py
@@ -150,6 +150,20 @@ class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2):
       * ``sparse_index_dim`` — width of the index-K/V vectors.
     """
 
+    # The AttentionOp-facing tensors this manager builds are synthetic
+    # placeholders over INDEX_KEY-coalesced pools, so one-model speculative
+    # draft layers must live in a separate manager even under attention DP
+    # (read by ``_should_create_separate_draft_kv_cache``).
+    supports_shared_draft_layers = False
+
+    # Workaround: the Eagle3 drafter's trtllm-gen generation kernels hit an
+    # illegal memory access at tokens_per_block=128 (base regression), while
+    # the MSA target requires 128. The dense draft layers have no page-size
+    # constraint, so the separate draft manager runs at 32 (read by
+    # ``_create_one_model_draft_kv_cache_manager``). Remove once the kernel
+    # bug is fixed.
+    draft_manager_tokens_per_block = 32
+
     def __init__(
         self,
         *args,
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py
index 86cd67fa7272..9e79fde9ccd0 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/msa_backend.py
@@ -246,6 +246,11 @@ def __post_init__(self) -> None:
         super().__post_init__()
         params = self.sparse_metadata_params
         self._msa_params = params if isinstance(params, MiniMaxM3SparseMetadataParams) else None
+        # See on_update_kv_lens.
+        self._msa_live_batch = 0
+        self._msa_live_total_q = 0
+        self._msa_page_size = 0
+        self._msa_corrected_kv_lens_cpu: Optional[torch.Tensor] = None
         self._create_msa_buffers()
 
     @property
@@ -259,11 +264,23 @@ def msa_qo_lens_cpu(self) -> Optional[torch.Tensor]:
 
     @property
     def msa_kv_lens_cpu(self) -> Optional[torch.Tensor]:
-        """Per-request KV length, cached plus new tokens (host int32)."""
+        """Per-request KV length, cached plus new tokens (host int32).
+
+        The base ``kv_lens`` includes ``num_extra_kv_tokens`` (speculative
+        draft-loop slots consumed by the C++ kernels); the MSA plans, ladder
+        slots and page counts need the true attended length, so it is
+        excluded here.
+        """
+        if self._msa_corrected_kv_lens_cpu is not None:
+            return self._msa_corrected_kv_lens_cpu
         kv_lens = getattr(self, "kv_lens", None)
         if self.seq_lens is None or kv_lens is None:
             return None
         out = kv_lens[: self.num_seqs]
+        params = self.kv_cache_params
+        extra = params.num_extra_kv_tokens if params is not None else 0
+        if extra:
+            out = out - extra
         return out if out.dtype == torch.int32 else out.to(torch.int32)
 
     @property
@@ -326,6 +343,37 @@ def _create_msa_buffers(self) -> None:
             dtype=torch.int32,
             capture_graph=capture_graph,
         )
+        # Staging for on_update_kv_lens: re-derives slots/bounds on device
+        # from the corrected kv_lens_cuda, sync-free.
+        tokens_per_block = int(kv_cache_manager.tokens_per_block)
+        self.msa_req_to_token = self.get_empty(
+            buffers,
+            (max_num_sequences, max_blocks_per_seq * tokens_per_block),
+            cache_name="msa_req_to_token",
+            dtype=torch.int32,
+            capture_graph=capture_graph,
+        )
+        self.msa_q_batch_row = self.get_empty(
+            buffers,
+            (max_num_tokens,),
+            cache_name="msa_q_batch_row",
+            dtype=torch.int32,
+            capture_graph=capture_graph,
+        )
+        self.msa_q_intra = self.get_empty(
+            buffers,
+            (max_num_tokens,),
+            cache_name="msa_q_intra",
+            dtype=torch.int32,
+            capture_graph=capture_graph,
+        )
+        self.msa_qo_lens_dev = self.get_empty(
+            buffers,
+            (max_num_sequences,),
+            cache_name="msa_qo_lens_dev",
+            dtype=torch.int32,
+            capture_graph=capture_graph,
+        )
         # The proxy scratch needs the fmha_sm100 plan geometry. This metadata
         # exists only for the MSA backend, whose selection already required the
         # kernels, so a failed import here is a hard error rather than a reason
@@ -341,37 +389,52 @@ def _create_msa_buffers(self) -> None:
             )
             self._alloc_msa_proxy_scratch(
                 num_index_heads=params.num_index_heads,
-                max_batch=max_num_sequences,
+                max_tokens=self._msa_max_decode_tokens(),
                 max_k_tiles=max_k_tiles,
                 capture_graph=capture_graph,
             )
         self._msa_buffers_ready = True
 
+    def _msa_max_decode_tokens(self) -> int:
+        """Worst-case decode-step query tokens (spec verify emits 1 + draft_len
+        per row), bounded by max_num_tokens. getattr fallbacks cover metadata
+        built via ``__new__`` in structural tests.
+        """
+        max_seqs = int(getattr(self, "max_num_sequences", 0) or 0)
+        max_toks = int(getattr(self, "max_num_tokens", 0) or 0)
+        if max_toks <= 0:
+            return max_seqs
+        # 16384 keeps total_q * num_qo_heads under the fmha_sm100 planner cap
+        # (65536, fmha_sm100/api.py) at 4 sharded index heads.
+        return max(max_seqs, min(max_toks, 16384))
+
     def _alloc_msa_proxy_scratch(
         self,
         *,
         num_index_heads: int,
-        max_batch: int,
+        max_tokens: int,
         max_k_tiles: int,
         capture_graph: bool,
     ) -> None:
         """Allocate the flat proxy max-score store and the valid-block scratch.
 
-        The store is sized for the worst-case max_k_tiles so one allocation
-        serves every decode step. msa_proxy_max_score_view slices the per-step
-        shape out of it.
+        The store is sized for the worst-case max_k_tiles and the worst-case
+        per-step query-token count (which exceeds the batch size under
+        speculative multi-token verify), so one allocation serves every
+        decode step. msa_proxy_max_score_view slices the per-step shape out
+        of it.
         """
         buffers = self.cuda_graph_buffers
         self.msa_max_score = self.get_empty(
             buffers,
-            (num_index_heads * max_k_tiles * max_batch,),
+            (num_index_heads * max_k_tiles * max_tokens,),
             cache_name="msa_max_score",
             dtype=torch.float32,
             capture_graph=capture_graph,
         )
         self.msa_n_valid_blocks = self.get_empty(
             buffers,
-            (max_batch,),
+            (max_tokens,),
             cache_name="msa_n_valid_blocks",
             dtype=torch.int32,
             capture_graph=capture_graph,
@@ -386,14 +449,15 @@ def _ensure_msa_decode_scratch_buffers(
         required_max_k_tiles: int,
     ) -> None:
         """Ensure proxy scratch buffers exist and cover the current plan."""
-        required_numel = num_index_heads * required_max_k_tiles * max_batch
+        max_tokens = max(int(max_batch), self._msa_max_decode_tokens())
+        required_numel = num_index_heads * required_max_k_tiles * max_tokens
         if self.msa_max_score is not None:
             if self.msa_max_score.numel() < required_numel:
                 raise ValueError(
                     f"msa_max_score backing store ({self.msa_max_score.numel()} "
                     f"elements) is smaller than the decode plan needs "
                     f"({required_numel} = {num_index_heads} heads * "
-                    f"{required_max_k_tiles} k-tiles * {max_batch} batch)."
+                    f"{required_max_k_tiles} k-tiles * {max_tokens} tokens)."
                 )
             return
 
@@ -415,16 +479,89 @@ def _ensure_msa_decode_scratch_buffers(
             )
         self._alloc_msa_proxy_scratch(
             num_index_heads=num_index_heads,
-            max_batch=max_batch,
+            max_tokens=max_tokens,
             max_k_tiles=max_k_tiles,
             capture_graph=capture_graph,
         )
 
     def prepare(self) -> None:
         super().prepare()
+        self._msa_corrected_kv_lens_cpu = None
         self._build_msa_fields()
         self._build_decode_plans()
 
+    def on_update_kv_lens(self) -> None:
+        """Re-derive length-dependent MSA state from the corrected kv_lens_cuda.
+
+        The overlap scheduler corrects kv_lens_cuda on device after prepare()
+        staged optimistic (full-acceptance) lens. Decode steps are patched
+        with pure device ops (capture-safe); the correction only shrinks
+        lengths, so the host-baked plan worklists and the page-table layout
+        stay valid. Eager mixed batches install corrected host lens and
+        rebuild (small D2H sync). Idempotent.
+        """
+        super().on_update_kv_lens()
+        if not self._msa_fields_ready:
+            return
+        batch = self._msa_live_batch
+        total_q = self._msa_live_total_q
+        if batch <= 0 or total_q <= 0:
+            return
+        if self.msa_decode_proxy_plan is None:
+            # Prefill/mixed step: the page-table layout and the host lens
+            # mirrors were built from the optimistic values, so install the
+            # corrected lens and rebuild; the per-layer eager plans read the
+            # corrected msa_kv_lens_cpu. Mixed steps are never captured.
+            if torch.cuda.is_current_stream_capturing():
+                return
+            self._msa_corrected_kv_lens_cpu = self.kv_lens_cuda[:batch].to("cpu", torch.int32)
+            self._build_msa_fields()
+            self._build_decode_plans()
+            return
+
+        kv_true = self.kv_lens_cuda[:batch]
+        qbr = self.msa_q_batch_row[:total_q].to(torch.long)
+        qo_dev = self.msa_qo_lens_dev[:batch]
+        kv_true_tok = kv_true[qbr]
+        pos = kv_true_tok - qo_dev[qbr] + self.msa_q_intra[:total_q]
+
+        # KV/idx-K write slots: slot[j] = req_to_token[request[j], pos[j]].
+        width = int(self.msa_req_to_token.shape[1])
+        idx = pos.to(torch.long).clamp(min=0, max=width - 1)
+        slots = self.msa_req_to_token.reshape(-1).index_select(0, qbr * width + idx)
+        self.msa_out_cache_loc[:total_q].copy_(slots)
+
+        # Per-token valid-block counts for top-k selection. clamp_min(1)
+        # keeps degenerate (graph-padding) rows from -inf-masking every
+        # block, which would NaN the fully-masked GQA row.
+        page = self._msa_page_size
+        n_valid = torch.div((pos + 1).clamp_min(1) + (page - 1), page, rounding_mode="floor")
+        self.msa_n_valid_blocks[:total_q].copy_(n_valid.to(torch.int32))
+
+        # Plan mirrors: proxy and dense keep one row per request; the sparse
+        # GQA plan is row-expanded per token with the ladder in the per-row
+        # offset. qo_offset must stay non-negative: negative values hit the
+        # kernel's packed-length sentinel fallback.
+        off_req = (kv_true - qo_dev).clamp_min(0)
+        for owner, expanded in (
+            (self._msa_proxy_plan, False),
+            (self._msa_gqa_plan, True),
+            (self._msa_dense_plan, False),
+        ):
+            decode = owner.plan[3]
+            seg_lens = decode.get("kv_segment_lens")
+            qo_off = decode.get("qo_offset")
+            if expanded:
+                if seg_lens is not None:
+                    seg_lens[:total_q].copy_(kv_true_tok.to(seg_lens.dtype))
+                if qo_off is not None:
+                    qo_off[:total_q].copy_(pos.clamp_min(0).to(qo_off.dtype))
+            else:
+                if seg_lens is not None:
+                    seg_lens[:batch].copy_(kv_true.to(seg_lens.dtype))
+                if qo_off is not None:
+                    qo_off[:batch].copy_(off_req.to(qo_off.dtype))
+
     def _build_decode_plans(self) -> None:
         """Build the graph-safe decode plans and buffers for this step.
 
@@ -458,7 +595,6 @@ def _build_decode_plans(self) -> None:
         qo_offset_cpu = self.msa_qo_offset_cpu
         if qo_lens_cpu is None or kv_lens_cpu is None or qo_offset_cpu is None:
             return
-        batch = int(qo_lens_cpu.shape[0])
         device = _cache_device(self)
         page_size = int(self.kv_cache_manager.tokens_per_block)
         capture_graph = self.is_cuda_graph
@@ -510,27 +646,31 @@ def _build_decode_plans(self) -> None:
         )
 
         # Allocate the graph-safe plan owners once per metadata; later steps
-        # only refresh their contents below.
+        # only refresh their contents below. The plan worklists are sized per
+        # expanded row (the planner splits qo_len > 1 requests into per-token
+        # rows under speculative multi-token verify), so use the worst-case
+        # decode-step token count rather than the batch size.
         if self._msa_proxy_plan is None:
+            max_plan_rows = max(max_batch, self._msa_max_decode_tokens())
             num_ctas = torch.cuda.get_device_properties(device).multi_processor_count
             self._msa_proxy_plan = _MsaGraphSafePlan(
                 self,
                 "msa_proxy_plan",
-                max_batch=max_batch,
+                max_batch=max_plan_rows,
                 num_ctas=num_ctas,
                 capture_graph=capture_graph,
             )
             self._msa_gqa_plan = _MsaGraphSafePlan(
                 self,
                 "msa_gqa_plan",
-                max_batch=max_batch,
+                max_batch=max_plan_rows,
                 num_ctas=num_ctas,
                 capture_graph=capture_graph,
             )
             self._msa_dense_plan = _MsaGraphSafePlan(
                 self,
                 "msa_dense_plan",
-                max_batch=max_batch,
+                max_batch=max_plan_rows,
                 num_ctas=num_ctas,
                 capture_graph=capture_graph,
             )
@@ -544,7 +684,9 @@ def _build_decode_plans(self) -> None:
         n_valid = per_token_valid_blocks(
             qo_lens_cpu, kv_lens_cpu, qo_offset_cpu, causal=True, block_size=page_size
         )
-        self.msa_n_valid_blocks[:batch].copy_(n_valid.to(torch.int32), non_blocking=True)
+        # One entry per query token (qo_len > 1 under spec verify).
+        total_q = int(n_valid.shape[0])
+        self.msa_n_valid_blocks[:total_q].copy_(n_valid.to(torch.int32), non_blocking=True)
 
     def _build_msa_fields(self) -> None:
         """Populate the MSA cache-write buffers for this step.
@@ -570,14 +712,6 @@ def _build_msa_fields(self) -> None:
         cache_device = _cache_device(self)
         page_size = int(kv_cache_manager.tokens_per_block)
 
-        is_prefill = int(self.num_contexts or 0) > 0
-        if not is_prefill and int(qo_lens_cpu.max().item()) > 1:
-            raise NotImplementedError(
-                "MiniMax-M3 MSA attention does not support speculative decoding "
-                "(multiple query tokens per decode step). Disable speculative "
-                "decoding or use the non-MSA MiniMax-M3 backend."
-            )
-
         # Built in prepare() (outside capture), so these transients are
         # fine: forwards read only the persistent buffers filled below.
         # qo_offset is the prefix length, so one build covers prefill
@@ -606,6 +740,25 @@ def _build_msa_fields(self) -> None:
 
         self.msa_out_cache_loc[:total_new_tokens].copy_(out_cache_loc, non_blocking=True)
         self.msa_kv_indices[:total_pages].copy_(kv_indices, non_blocking=True)
+
+        # Staging for on_update_kv_lens.
+        step_width = int(req_to_token.shape[1])
+        self.msa_req_to_token[:batch_size, :step_width].copy_(req_to_token, non_blocking=True)
+        qo_long = qo_lens_cpu.to(torch.long)
+        batch_row_cpu = torch.repeat_interleave(
+            torch.arange(batch_size, dtype=torch.int32), qo_long
+        )
+        starts = torch.cumsum(qo_long, 0) - qo_long
+        intra_cpu = (
+            torch.arange(total_new_tokens, dtype=torch.int64)
+            - torch.repeat_interleave(starts, qo_long)
+        ).to(torch.int32)
+        self.msa_q_batch_row[:total_new_tokens].copy_(batch_row_cpu)
+        self.msa_q_intra[:total_new_tokens].copy_(intra_cpu)
+        self.msa_qo_lens_dev[:batch_size].copy_(qo_lens_cpu)
+        self._msa_live_batch = batch_size
+        self._msa_live_total_q = total_new_tokens
+        self._msa_page_size = page_size
         self._msa_fields_ready = True
 
     def msa_idx_k_cache(self, layer_idx: int) -> torch.Tensor:
diff --git a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_metadata.py b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_metadata.py
index be4b8ae24143..a57cf5eac2e1 100644
--- a/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_metadata.py
+++ b/tensorrt_llm/_torch/attention_backend/sparse/minimax_m3/triton_metadata.py
@@ -20,7 +20,7 @@
 
 import torch
 
-from ...interface import AttentionMetadata
+from ...trtllm import TrtllmAttentionMetadata
 from .common import build_paged_kv_slot_mapping
 
 
@@ -93,6 +93,10 @@ class MiniMaxM3TritonSparseAttentionMetadata:
     prefix_lens: Optional[torch.Tensor] = None
     cu_seqlens_q: Optional[torch.Tensor] = None
     extend_seq_lens_cpu: Optional[List[int]] = None
+    # Query tokens per request on decode-shaped metadata: 1 normally,
+    # 1 + draft_len under one-model Eagle3 verify. Consumed by the dense
+    # SDPA decode ladder mask.
+    decode_qo_len: int = 1
     q_batch_row: Optional[torch.Tensor] = None
     q_positions: Optional[torch.Tensor] = None
     max_seqlen_q: int = field(default=1)
@@ -165,7 +169,11 @@ def prepare(self) -> None:
         if batch_size == 0:
             self.max_seqlen_k = 1
         else:
-            self.max_seqlen_k = int(self.seq_lens_cpu[:batch_size].max().item())
+            max_k = int(self.seq_lens_cpu[:batch_size].max().item())
+            # Optimistic overlap+spec lengths can overhang the page table at
+            # a page boundary; SDPA consumes max_seqlen_k as the exact
+            # mask/gather width, so clamp to the table.
+            self.max_seqlen_k = min(max_k, int(self.req_to_token.shape[1]))
 
 
 def ensure_metadata_on_device(
@@ -369,6 +377,44 @@ def _build_runtime_metadata_fresh(
     return meta, out_cache_loc
 
 
+def derive_q_positions_and_cache_slots(
+    req_to_token: torch.Tensor,
+    prefix_lens: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    q_batch_row: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Per-Q-token K-side positions and KV slot ids, on device, sync-free.
+
+    Shared by the metadata builder and the ``on_update_kv_lens``
+    re-derivation so the two cannot drift.
+    """
+    total_q = int(q_batch_row.shape[0])
+    qbr = q_batch_row.to(torch.long)
+    tok = torch.arange(total_q, dtype=torch.int32, device=q_batch_row.device)
+    q_positions = prefix_lens[qbr] + (tok - cu_seqlens_q[qbr])
+    # Optimistic prefix_lens can overhang the last allocated page; the
+    # overhanging slots are placeholders (on_update_kv_lens re-derives them
+    # before any forward reads them) but the gather must stay in bounds.
+    # Clamp only the table index — non-inplace, since ``.to`` aliases int64
+    # inputs.
+    idx = q_positions.to(torch.long).clamp(min=0, max=req_to_token.shape[1] - 1)
+    flat = qbr * req_to_token.shape[1] + idx
+    return q_positions, req_to_token.reshape(-1).index_select(0, flat)
+
+
+def derive_decode_cache_slots(req_to_token: torch.Tensor, seq_lens: torch.Tensor) -> torch.Tensor:
+    """Decode-row KV slot ids (new token at position ``seq_lens[b] - 1``).
+
+    Same in-bounds-placeholder clamp contract as
+    :func:`derive_q_positions_and_cache_slots`; the ``min=0`` floor
+    additionally covers zero-length dummy rows indexing ``-1``.
+    """
+    rows = torch.arange(seq_lens.shape[0], device=seq_lens.device, dtype=torch.long)
+    idx = (seq_lens.to(torch.long) - 1).clamp_(min=0, max=req_to_token.shape[1] - 1)
+    flat = rows * req_to_token.shape[1] + idx
+    return req_to_token.reshape(-1).index_select(0, flat)
+
+
 def build_runtime_metadata_from_kv_manager(
     *,
     kv_cache_manager,
@@ -539,21 +585,13 @@ def build_runtime_metadata_from_kv_manager(
         req_to_token = req_to_token_fresh
         slot_ids = torch.arange(batch, device=device, dtype=torch.int32)
 
-    # Compute out_cache_loc: per-new-token slot ids, in flattened order
-    # matching the q-token order the model layer projects. The Python
-    # loops below run on CPU lists derived from the CPU-resident
-    # ``seq_lens_cpu`` / ``prefix_lens`` / ``extend_seq_lens_cpu``, so
-    # no GPU sync is needed at this point. The resulting
-    # ``out_cache_loc`` tensor is constructed directly on ``device``.
-    # The ``int(...item())`` reads against ``req_to_token`` are a CPU
-    # sync but only ever run from ``prepare()`` (outside any CUDA-graph
-    # capture window) — they are not in the forward path.
+    # out_cache_loc must be flattened in the q-token order the model layer
+    # projects, or K/V lands in the wrong requests' slots.
     if is_prefill:
         if extend_seq_lens_cpu is None:
             raise ValueError("prefill metadata requires extend_seq_lens_cpu")
         if prefix_lens is None:
             raise ValueError("prefill metadata requires prefix_lens")
-        prefix_lens_cpu = prefix_lens.to("cpu").tolist()
         if static_buffers is not None:
             prefix_buf = static_buffers["prefix_lens"]
             prefix_src = prefix_lens.to(device=device, dtype=torch.int32, non_blocking=True)
@@ -563,17 +601,18 @@ def build_runtime_metadata_from_kv_manager(
             prefix_lens_dev = (
                 prefix_lens.to(device) if prefix_lens.device != device else prefix_lens
             )
-        out_cache_loc_list: List[int] = []
         cu_q: List[int] = [0]
-        req_to_token_cpu = req_to_token_fresh.to("cpu")
-        for b in range(batch):
-            pref = int(prefix_lens_cpu[b])
-            ext = int(extend_seq_lens_cpu[b])
-            for offset in range(ext):
-                slot = int(req_to_token_cpu[b, pref + offset].item())
-                out_cache_loc_list.append(slot)
-            cu_q.append(cu_q[-1] + ext)
+        for ext in extend_seq_lens_cpu:
+            cu_q.append(cu_q[-1] + int(ext))
         total_q = cu_q[-1]
+        cu_seqlens_q_src = torch.tensor(cu_q, dtype=torch.int32, device=device)
+        q_batch_row_src = torch.repeat_interleave(
+            torch.arange(batch, device=device, dtype=torch.int32),
+            torch.tensor(extend_seq_lens_cpu, dtype=torch.int64, device=device),
+        )
+        q_positions_src, out_cache_loc_src = derive_q_positions_and_cache_slots(
+            req_to_token, prefix_lens_dev, cu_seqlens_q_src, q_batch_row_src
+        )
         if static_buffers is not None:
             if total_q > static_buffers["max_num_tokens"]:
                 raise ValueError(
@@ -581,33 +620,22 @@ def build_runtime_metadata_from_kv_manager(
                     f"is smaller than current total_q={total_q}"
                 )
             out_cache_loc_buf = static_buffers["out_cache_loc"]
-            out_cache_loc_src = torch.tensor(out_cache_loc_list, dtype=torch.int32, device=device)
             out_cache_loc_buf[:total_q].copy_(out_cache_loc_src, non_blocking=True)
             out_cache_loc = out_cache_loc_buf[:total_q]
             cu_seqlens_q_buf = static_buffers["cu_seqlens_q"]
-            cu_seqlens_q_src = torch.tensor(cu_q, dtype=torch.int32, device=device)
             cu_seqlens_q_buf[: batch + 1].copy_(cu_seqlens_q_src, non_blocking=True)
             cu_seqlens_q = cu_seqlens_q_buf[: batch + 1]
-            # Populate persistent q_batch_row / q_positions in-place so
-            # the inner metadata's prepare() can leave them alone.
             q_batch_row_buf = static_buffers["q_batch_row"]
             q_positions_buf = static_buffers["q_positions"]
-            for b in range(batch):
-                start, end = cu_q[b], cu_q[b + 1]
-                if end > start:
-                    q_batch_row_buf[start:end] = b
-                    pref = int(prefix_lens_cpu[b])
-                    offsets = (
-                        torch.arange(start, end, device=device, dtype=torch.int32) - start + pref
-                    )
-                    q_positions_buf[start:end].copy_(offsets, non_blocking=True)
+            q_batch_row_buf[:total_q].copy_(q_batch_row_src, non_blocking=True)
+            q_positions_buf[:total_q].copy_(q_positions_src, non_blocking=True)
             q_batch_row = q_batch_row_buf[:total_q]
             q_positions = q_positions_buf[:total_q]
         else:
-            out_cache_loc = torch.tensor(out_cache_loc_list, dtype=torch.int32, device=device)
-            cu_seqlens_q = torch.tensor(cu_q, dtype=torch.int32, device=device)
-            q_batch_row = None
-            q_positions = None
+            out_cache_loc = out_cache_loc_src
+            cu_seqlens_q = cu_seqlens_q_src
+            q_batch_row = q_batch_row_src
+            q_positions = q_positions_src
         meta = MiniMaxM3TritonSparseAttentionMetadata(
             is_prefill=True,
             req_to_token=req_to_token,
@@ -622,12 +650,7 @@ def build_runtime_metadata_from_kv_manager(
         )
     else:
         # Decode: the new token sits at position seq_lens[b] - 1.
-        seq_lens_cpu_list = seq_lens_cpu.to("cpu").tolist()
-        out_cache_loc_list = []
-        req_to_token_cpu = req_to_token_fresh.to("cpu")
-        for b in range(batch):
-            pos = int(seq_lens_cpu_list[b]) - 1
-            out_cache_loc_list.append(int(req_to_token_cpu[b, pos].item()))
+        out_cache_loc_src = derive_decode_cache_slots(req_to_token, seq_lens_dev)
         if static_buffers is not None:
             if batch > static_buffers["max_num_tokens"]:
                 raise ValueError(
@@ -635,11 +658,10 @@ def build_runtime_metadata_from_kv_manager(
                     f"is smaller than current batch={batch}"
                 )
             out_cache_loc_buf = static_buffers["out_cache_loc"]
-            out_cache_loc_src = torch.tensor(out_cache_loc_list, dtype=torch.int32, device=device)
             out_cache_loc_buf[:batch].copy_(out_cache_loc_src, non_blocking=True)
             out_cache_loc = out_cache_loc_buf[:batch]
         else:
-            out_cache_loc = torch.tensor(out_cache_loc_list, dtype=torch.int32, device=device)
+            out_cache_loc = out_cache_loc_src
         meta = MiniMaxM3TritonSparseAttentionMetadata(
             is_prefill=False,
             req_to_token=req_to_token,
@@ -651,8 +673,13 @@ def build_runtime_metadata_from_kv_manager(
     return meta, out_cache_loc
 
 
-class MiniMaxM3AttentionMetadata(AttentionMetadata):
-    """:class:`AttentionMetadata` that pre-builds MiniMax-M3 metadata.
+class MiniMaxM3AttentionMetadata(TrtllmAttentionMetadata):
+    """:class:`TrtllmAttentionMetadata` that pre-builds MiniMax-M3 metadata.
+
+    Subclasses :class:`TrtllmAttentionMetadata` (precedent:
+    ``DSAtrtllmAttentionMetadata``): one-model Eagle3 draft layers run
+    :class:`TrtllmAttention` against this shared per-step metadata, and the
+    engine gates spec-dec plumbing on the TRTLLM ``isinstance``.
 
     Overrides :meth:`prepare` so the M3-sparse
     :class:`MiniMaxM3TritonSparseAttentionMetadata` and the per-new-token
@@ -824,7 +851,9 @@ def prepare(self) -> None:
         # specialization. (iter-131 regression: previously a wrong
         # predicate routed mixed batches into the decode branch and
         # crashed in index_copy_.)
-        is_extend = num_contexts > 0
+        # Multi-token gen rows (spec verify) also route through the extend
+        # path as prefix+window extends; decode stays one-token-per-row.
+        is_extend = num_contexts > 0 or int(seq_lens_cpu[:batch_size].max().item()) > 1
         if is_extend:
             prefix_lens_list = [int(num_cached_per_seq[b]) for b in range(batch_size)]
             extend_seq_lens_cpu = [
@@ -862,11 +891,52 @@ def prepare(self) -> None:
             "out_cache_loc": out_cache_loc,
         }
 
+    def on_update_kv_lens(self) -> None:
+        """Re-derive the M3 attachment from the corrected ``kv_lens_cuda``.
+
+        Under the overlap scheduler + speculative decoding, prepare()
+        runs with optimistic cached counts (full draft acceptance) and
+        the engine corrects ``kv_lens_cuda`` on device before invoking
+        this hook (same pattern as ``DSAtrtllmAttentionMetadata``).
+        On-device, sync-free, and idempotent; ``seq_lens_cpu`` /
+        ``max_seqlen_k`` keep the optimistic values — they only bound
+        arange widths that the kernels mask by ``seq_lens``.
+        """
+        super().on_update_kv_lens()
+        attachment = self.minimax_m3
+        if not attachment:
+            return
+        meta = attachment["metadata"]
+        out_cache_loc = attachment["out_cache_loc"]
+        batch = int(meta.slot_ids.shape[0])
+        kv_lens = self.kv_lens_cuda[:batch]
+        meta.seq_lens[:batch].copy_(kv_lens)
+        if meta.is_prefill:
+            # Only the K-side prefix moves with rejections; the Q-side
+            # structure (cu_seqlens_q, q_batch_row) is fixed per step.
+            total_q = int(meta.q_positions.shape[0])
+            cu = meta.cu_seqlens_q
+            meta.prefix_lens[:batch].copy_(kv_lens - (cu[1 : batch + 1] - cu[:batch]))
+            q_positions, cache_slots = derive_q_positions_and_cache_slots(
+                meta.req_to_token,
+                meta.prefix_lens[:batch],
+                cu,
+                meta.q_batch_row[:total_q],
+            )
+            meta.q_positions[:total_q].copy_(q_positions)
+            out_cache_loc[:total_q].copy_(cache_slots)
+        else:
+            # Identity today; keeps 0-draft corrections right if they
+            # become reachable.
+            out_cache_loc[:batch].copy_(derive_decode_cache_slots(meta.req_to_token, kv_lens))
+
 
 __all__ = [
     "MiniMaxM3AttentionMetadata",
     "MiniMaxM3TritonSparseAttentionMetadata",
     "allocate_minimax_m3_static_buffers",
     "build_runtime_metadata_from_kv_manager",
+    "derive_decode_cache_slots",
+    "derive_q_positions_and_cache_slots",
     "ensure_metadata_on_device",
 ]
diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py
index f10696c97bb9..2d6945ae06a2 100644
--- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py
+++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py
@@ -55,6 +55,7 @@
 from ..modules.linear import Linear, TensorParallelMode, copy_weight, load_weight_shard
 from ..modules.multi_stream_utils import maybe_execute_in_parallel
 from ..modules.rms_norm import RMSNorm
+from ..speculative import SpecMetadata
 from ..utils import (
     ActivationType,
     AuxStreamType,
@@ -64,7 +65,8 @@
 )
 from .checkpoints.base_weight_mapper import BaseWeightMapper
 from .checkpoints.hf.minimaxm3_weight_mapper import MINIMAX_M3_PARAMS_MAP, MiniMaxM3HfWeightMapper
-from .modeling_utils import DecoderModel, DecoderModelForCausalLM, ModelConfig, register_auto_model
+from .modeling_speculative import SpecDecOneEngineForCausalLM
+from .modeling_utils import DecoderModel, ModelConfig, register_auto_model
 
 # Dense layers use SDPA with non-contiguous Q/K/V and a bool attn_mask.
 # Limit backends to memory-efficient and math; cuDNN SDPA fails for this layout,
@@ -974,7 +976,17 @@ def _dense_attention_core(
 
         # 7. Gather padded K/V for every batch row and run dense GQA.
         batch = int(m3_meta.slot_ids.shape[0])
-        max_k = int(m3_meta.max_seqlen_k)
+        # Graph capture bakes the gather/mask width, and max_seqlen_k is a
+        # per-step host bound later replays can outgrow. Bake
+        # min(page-table width, engine max_seq_len): the raw table width
+        # alone is inflated far past max_seq_len by the KV-estimation pass
+        # and would OOM the [batch, max_k, heads] gather. The seq_lens mask
+        # below invalidates the slack.
+        if attn_metadata.is_cuda_graph:
+            capacity = int(m3_meta.req_to_token.shape[1])
+            max_k = min(capacity, int(attn_metadata.max_seq_len or capacity))
+        else:
+            max_k = int(m3_meta.max_seqlen_k)
         if max_k <= 0:
             max_k = 1
         # ``_gather_paged_batched`` decomposes the flat slot id into
@@ -1001,15 +1013,11 @@ def _dense_attention_core(
                 f"by num_key_value_heads ({self.num_key_value_heads})"
             )
         group = self.num_heads // max(self.num_key_value_heads, 1)
-        if group > 1:
-            k_padded = k_padded.repeat_interleave(group, dim=2)
-            v_padded = v_padded.repeat_interleave(group, dim=2)
 
         # Build the per-query attention mask that masks out padded KV
         # positions beyond each sequence's true ``seq_lens`` and (for
         # prefill) preserves causality. ``q_positions`` from the prefill
-        # metadata names each Q token's K-side position; for decode
-        # there is one Q token per request at position ``seq_lens - 1``.
+        # metadata names each Q token's K-side position.
         # The metadata tensors are produced by
         # :meth:`MiniMaxM3AttentionMetadata.prepare` on the cache
         # device, so ``.to(dtype=torch.long)`` is a same-device dtype
@@ -1018,6 +1026,9 @@ def _dense_attention_core(
         kv_positions = torch.arange(max_k, device=q.device).unsqueeze(0)  # [1, max_k]
 
         if m3_meta.is_prefill:
+            if group > 1:
+                k_padded = k_padded.repeat_interleave(group, dim=2)
+                v_padded = v_padded.repeat_interleave(group, dim=2)
             # Prefill: build [total_q, max_k] mask using q_positions / q_batch_row.
             # Prefill never runs inside the CUDA-graph capture window
             # (capture is decode-only), so the per-batch Python loop and
@@ -1058,35 +1069,42 @@ def _dense_attention_core(
                     )  # [1, H, q, d]
                 output_view[start:end].copy_(out_b.squeeze(0).transpose(0, 1))
         else:
-            # Decode: one Q token per request at position seq_lens - 1.
-            # Every input tensor here is already on q.device (set up by
-            # prepare()), so SDPA captures cleanly.
-            valid = kv_positions < seq_lens_dev.unsqueeze(-1)  # [batch, max_k]
-            q_b = q_view.unsqueeze(1).transpose(1, 2)  # [batch, H, 1, d]
-            k_b = k_padded.transpose(1, 2)  # [batch, H, k, d]
-            v_b = v_padded.transpose(1, 2)  # [batch, H, k, d]
-            mask_b = valid.unsqueeze(1).unsqueeze(1)  # [batch, 1, 1, k]
+            # Decode: qo_len query tokens per request; token t of request b
+            # attends seq_lens[b] - qo_len + t + 1 positions (the causal
+            # ladder; qo_len=1 is the classic one-token mask). Every input
+            # tensor here is already on q.device (set up by prepare()), so
+            # SDPA captures cleanly.
+            qo_len = int(m3_meta.decode_qo_len)
+            ladder = torch.arange(1 - qo_len, 1, device=q.device, dtype=torch.long)
+            # eff[b, t] = attendable position count for token t of row b.
+            eff = seq_lens_dev.unsqueeze(-1) + ladder  # [batch, qo]
+            valid = kv_positions.unsqueeze(1) < eff.unsqueeze(-1)  # [batch, qo, max_k]
+            q_b = q_view.view(batch, qo_len, self.num_heads, self.head_dim).transpose(
+                1, 2
+            )  # [batch, H, qo, d]
+            mask_b = valid.unsqueeze(1)  # [batch, 1, qo, k]
+            # Expand K/V one KV head at a time: the all-heads transient is
+            # O(batch * max_k * num_heads) inside the CUDA-graph pool and
+            # exceeds the pool budget under attention DP (unsharded heads);
+            # with TP-sharded KV heads this is one iteration.
+            out_b = q.new_empty(batch, self.num_heads, qo_len, self.head_dim)
             with sdpa_kernel(_DENSE_SDPA_BACKENDS):
-                out_b = torch.nn.functional.scaled_dot_product_attention(
-                    q_b.to(q.dtype),
-                    k_b.to(q.dtype),
-                    v_b.to(q.dtype),
-                    attn_mask=mask_b,
-                    dropout_p=0.0,
-                    is_causal=False,
-                )  # [batch, H, 1, d]
-            # Drop the singleton Q-length axis and write the resulting
-            # ``[batch, num_heads, head_dim]`` tensor into the final buffer.
-            # The prior ``.transpose(1, 2).reshape(batch, H, d)`` pattern
-            # was wrong: with ``H != head_dim`` (M3 TP=8 has H=8, d=128)
-            # the non-contiguous transpose forces ``reshape`` to copy the
-            # data in C-order under its current ``[batch, d, H]`` shape,
-            # then reinterpret as ``[batch, H, d]`` — which scrambles
-            # ``(head, head_dim)`` ordering and feeds permuted activations
-            # into ``o_proj``. Prefill is unaffected because its
-            # ``transpose(0, 1)`` runs between q-len and num_heads axes
-            # which the per-batch loop already laid out correctly.
-            output.view(batch, self.num_heads, self.head_dim).copy_(out_b.squeeze(2))
+                for h in range(max(self.num_key_value_heads, 1)):
+                    qh = slice(h * group, (h + 1) * group)
+                    k_h = k_padded[:, :, h : h + 1].repeat_interleave(group, dim=2)
+                    v_h = v_padded[:, :, h : h + 1].repeat_interleave(group, dim=2)
+                    out_b[:, qh] = torch.nn.functional.scaled_dot_product_attention(
+                        q_b[:, qh].to(q.dtype),
+                        k_h.transpose(1, 2).to(q.dtype),
+                        v_h.transpose(1, 2).to(q.dtype),
+                        attn_mask=mask_b,
+                        dropout_p=0.0,
+                        is_causal=False,
+                    )  # [batch, group, qo, d]
+            # Copy through a token-major [batch, qo, H, dh] view rather
+            # than transpose(1, 2).reshape, which (with H != head_dim)
+            # copies in C-order and scrambles (head, head_dim) into o_proj.
+            output.view(batch, qo_len, self.num_heads, self.head_dim).copy_(out_b.transpose(1, 2))
 
         return output
 
@@ -1381,6 +1399,7 @@ def forward(
         hidden_states: torch.Tensor,
         attn_metadata: AttentionMetadata,
         residual: Optional[torch.Tensor],
+        spec_metadata: Optional[SpecMetadata] = None,
         **kwargs,
     ) -> torch.Tensor:
         if residual is None:
@@ -1401,6 +1420,10 @@ def forward(
             hidden_states = self.block_sparse_moe(hidden_states, attn_metadata)
         else:
             hidden_states = self.mlp(hidden_states)
+        # hidden_states is fully TP-reduced at layer exit (no cross-layer
+        # allreduce+norm fusion).
+        if spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx):
+            spec_metadata.maybe_capture_hidden_states(self.layer_idx, hidden_states, residual)
         return hidden_states, residual
 
 
@@ -1451,6 +1474,7 @@ def forward(
         input_ids: Optional[torch.IntTensor] = None,
         position_ids: Optional[torch.IntTensor] = None,
         inputs_embeds: Optional[torch.FloatTensor] = None,
+        spec_metadata: Optional[SpecMetadata] = None,
         **kwargs,
     ) -> torch.Tensor:
         if (input_ids is None) ^ (inputs_embeds is not None):
@@ -1467,6 +1491,7 @@ def forward(
                 hidden_states=hidden_states,
                 attn_metadata=attn_metadata,
                 residual=residual,
+                spec_metadata=spec_metadata,
             )
 
         hidden_states, _ = self.norm(hidden_states, residual)
@@ -1474,19 +1499,14 @@ def forward(
 
 
 @register_auto_model("MiniMaxM3SparseForCausalLM")
-class MiniMaxM3ForCausalLM(DecoderModelForCausalLM[MiniMaxM3Model, PretrainedConfig]):
+class MiniMaxM3ForCausalLM(SpecDecOneEngineForCausalLM[MiniMaxM3Model, PretrainedConfig]):
     """Text-only M3 model."""
 
     def __init__(self, model_config: "ModelConfig[PretrainedConfig]"):
         raw_pretrained = model_config.pretrained_config
         if is_minimax_m3_vl_config(raw_pretrained):
             model_config = get_text_model_config(model_config)
-        super().__init__(
-            MiniMaxM3Model(model_config),
-            config=model_config,
-            hidden_size=model_config.pretrained_config.hidden_size,
-            vocab_size=model_config.pretrained_config.vocab_size,
-        )
+        super().__init__(MiniMaxM3Model(model_config), model_config)
 
     def load_weights(
         self,
diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 0149ba7ae03a..db3eb0baaa5c 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -480,10 +480,8 @@ def _get_kv_size_per_token(self,
                 # External drafter: layers start from 0, normal PP distribution
                 # Resolve draft manager class from draft config — may differ
                 # from target (e.g. hybrid target + plain transformer draft).
-                draft_kv_cache_manager_cls = get_kv_cache_manager_cls(
-                    effective_draft_config,
-                    kv_cache_config,
-                    is_disagg=self._is_disagg)
+                draft_kv_cache_manager_cls = self._get_draft_kv_cache_manager_cls(
+                    effective_draft_config, kv_cache_config)
                 total += self._per_manager_cache_cost(
                     draft_kv_cache_manager_cls, effective_draft_config,
                     kv_cache_config)
@@ -1048,10 +1046,17 @@ def _should_create_separate_draft_kv_cache(self) -> bool:
         in the target model and don't produce a separate ModelConfig. We fall
         back to the target model's config via _get_effective_draft_config().
         """
-        if self._mapping.enable_attention_dp:
-            logger.info(
-                "Attention DP is enabled, separate draft KV cache is not supported."
-            )
+        if self._mapping.enable_attention_dp and getattr(
+                self._kv_cache_manager_cls, 'supports_shared_draft_layers',
+                True):
+            # Under attention DP, draft layers share the target manager (the
+            # layout existing deployments were validated with). A manager can
+            # opt out: MiniMax-M3's 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 requires the
+            # separate draft manager even under attention DP.
+            logger.info("Attention DP: draft layers share the target KV "
+                        "cache manager.")
             return False
         return should_use_separate_draft_kv_cache(self._speculative_config)
 
@@ -1082,6 +1087,17 @@ def _get_num_draft_layers(self) -> int:
             return self._draft_config.pretrained_config.num_hidden_layers
         return get_num_spec_layers(self._speculative_config)
 
+    def _get_draft_kv_cache_manager_cls(self, effective_draft_config,
+                                        draft_kv_config):
+        """Resolve the draft manager class from the draft config, promoted
+        to V2 when the target manager is V2."""
+        draft_kv_cache_manager_cls = get_kv_cache_manager_cls(
+            effective_draft_config, draft_kv_config, is_disagg=self._is_disagg)
+        if self._is_kv_cache_manager_v2 and not issubclass(
+                draft_kv_cache_manager_cls, KVCacheManagerV2):
+            draft_kv_cache_manager_cls = KVCacheManagerV2
+        return draft_kv_cache_manager_cls
+
     def _create_one_model_draft_kv_cache_manager(
         self,
         estimating_kv_cache: bool = False,
@@ -1146,8 +1162,8 @@ def _create_one_model_draft_kv_cache_manager(
                 f"Derived draft KV cache max_attention_window for separate "
                 f"draft manager: {draft_kv_config.max_attention_window}")
         # Get the appropriate KV cache manager class for the draft model
-        draft_kv_cache_manager_cls = get_kv_cache_manager_cls(
-            effective_draft_config, draft_kv_config, is_disagg=self._is_disagg)
+        draft_kv_cache_manager_cls = self._get_draft_kv_cache_manager_cls(
+            effective_draft_config, draft_kv_config)
         draft_kv_cache_manager_cls = self._fallback_if_unsupported_kv_cache_manager_v2(
             draft_kv_cache_manager_cls, effective_draft_config, draft_kv_config)
 
@@ -1157,12 +1173,23 @@ def _create_one_model_draft_kv_cache_manager(
         # the sparse_attention_config. Get it from effective_draft_config which
         # falls back to the target model's config for MTP mode.
         sparse_attn_config = effective_draft_config.sparse_attention_config
+        # A target manager class may request a different page size for the
+        # separate draft manager (e.g. MiniMax-M3, see
+        # draft_manager_tokens_per_block there for the rationale).
+        draft_tpb = getattr(self._kv_cache_manager_cls,
+                            'draft_manager_tokens_per_block',
+                            self._tokens_per_block)
+        if draft_tpb != self._tokens_per_block:
+            logger.info(
+                f"Draft KV cache manager uses tokens_per_block={draft_tpb} "
+                f"(target uses {self._tokens_per_block}).")
+            draft_kv_config.tokens_per_block = draft_tpb
         return _create_kv_cache_manager(
             model_engine=None,
             kv_cache_manager_cls=draft_kv_cache_manager_cls,
             mapping=self._mapping,
             kv_cache_config=draft_kv_config,
-            tokens_per_block=self._tokens_per_block,
+            tokens_per_block=draft_tpb,
             max_seq_len=self._max_seq_len,
             max_batch_size=self._max_batch_size,
             spec_config=self._speculative_config,
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index 34ebd340a5ee..e830afa4a7bd 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -5777,6 +5777,11 @@ def _pad_attention_dp_dummy_request(self):
                 and self.max_num_tokens is not None):
             token_nums = [self.max_num_tokens]
 
+        # A separate draft KV cache manager must also see the dummy, or
+        # its prepare_resources hits an unknown request id.
+        draft_kv_cache_manager = self.resource_manager.get_resource_manager(
+            ResourceManagerType.DRAFT_KV_CACHE_MANAGER)
+
         if (not self._enable_dsv4_adp_dummy_fixes
                 or self.kv_cache_transceiver is None):
             llm_request = self.kv_cache_manager.add_dummy_requests(
@@ -5785,6 +5790,7 @@ def _pad_attention_dp_dummy_request(self):
                 is_gen=self._adp_dummy_is_gen,
                 prepare_resource=True,
                 max_num_draft_tokens=self.max_total_draft_tokens,
+                draft_kv_cache_manager=draft_kv_cache_manager,
             )[0]
             llm_request.is_attention_dp_dummy = True
             spec_resource_manager = self.resource_manager.get_resource_manager(
@@ -5809,6 +5815,7 @@ def _pad_attention_dp_dummy_request(self):
                 is_gen=self._adp_dummy_is_gen,
                 prepare_resource=True,
                 max_num_draft_tokens=self.max_total_draft_tokens,
+                draft_kv_cache_manager=draft_kv_cache_manager,
             )
         except OutOfPagesError:
             dummy_requests = None
diff --git a/tensorrt_llm/_torch/speculative/eagle3.py b/tensorrt_llm/_torch/speculative/eagle3.py
index a1f4f82869bd..df8aebdb7ff8 100644
--- a/tensorrt_llm/_torch/speculative/eagle3.py
+++ b/tensorrt_llm/_torch/speculative/eagle3.py
@@ -658,8 +658,17 @@ def __init__(self,
     def max_draft_len(self) -> int:
         return self.spec_config.max_draft_len
 
-    def _prepare_attn_metadata_for_spec_dec(self, attn_metadata):
-        attn_metadata.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda")
+    def _prepare_attn_metadata_for_spec_dec(self, attn_metadata, spec_metadata):
+        # Graph warmup runs twice before capture while the draft loop
+        # mutates kv_lens_cuda in place — save/restore it during warmup
+        # (same as DFlash/PARD). During capture the mutation must be
+        # recorded, so skip the save there.
+        is_capturing = torch.cuda.is_current_stream_capturing()
+        if spec_metadata.is_cuda_graph and not is_capturing:
+            attn_metadata.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda",
+                                               "kv_lens_cuda")
+        else:
+            attn_metadata.prepare_for_spec_dec("_seq_lens", "_seq_lens_cuda")
         batch_size = attn_metadata.num_seqs
 
         # Save spec-dec params that the drafting loop will overwrite.
@@ -767,7 +776,7 @@ def _forward_impl(self,
             )
 
         # Save the old attn_metadata and spec_metadata
-        self._prepare_attn_metadata_for_spec_dec(attn_metadata)
+        self._prepare_attn_metadata_for_spec_dec(attn_metadata, spec_metadata)
 
         # Prepare inputs for the 1st draft model forward
         position_ids = position_ids.squeeze(0)
diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml
index ef32f74f1ba0..0f81e7f07335 100644
--- a/tests/integration/defs/accuracy/references/gsm8k.yaml
+++ b/tests/integration/defs/accuracy/references/gsm8k.yaml
@@ -495,6 +495,9 @@ nvidia/MiniMax-M3-NVFP4:
   - quant_algo: MIXED_PRECISION
     kv_cache_quant_algo: FP8
     accuracy: 86
+  - quant_algo: MIXED_PRECISION
+    spec_dec_algo: Eagle3
+    accuracy: 88
 nvidia/NVIDIA-Nemotron-Nano-9B-v2:
   - accuracy: 85.027
   - quant_algo: FP8
diff --git a/tests/integration/defs/accuracy/references/mmlu.yaml b/tests/integration/defs/accuracy/references/mmlu.yaml
index a870fbba6402..c332eb57e35b 100644
--- a/tests/integration/defs/accuracy/references/mmlu.yaml
+++ b/tests/integration/defs/accuracy/references/mmlu.yaml
@@ -287,6 +287,9 @@ nvidia/MiniMax-M3-NVFP4:
   - quant_algo: MIXED_PRECISION
     kv_cache_quant_algo: FP8
     accuracy: 81
+  - quant_algo: MIXED_PRECISION
+    spec_dec_algo: Eagle3
+    accuracy: 83
 moonshotai/Kimi-K2-Instruct:
   - quant_algo: FP8_BLOCK_SCALES
     accuracy: 87.65
diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
index 378aa74751db..9c408f9617a3 100644
--- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py
+++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
@@ -7846,6 +7846,105 @@ def test_nvfp4(self, use_msa):
             task = GSM8K(model_name)
             task.evaluate(llm)
 
+    @pytest.mark.skip_less_device(4)
+    @pytest.mark.skip_less_device_memory(140000)
+    @parametrize_with_ids("cuda_graph", [True])
+    @parametrize_with_ids("use_msa", [True])
+    @parametrize_with_ids("overlap_scheduler", [False, True])
+    @parametrize_with_ids("attention_dp", [False, True])
+    @parametrize_with_ids("tp_size,ep_size", [(4, 4)])
+    def test_nvfp4_eagle3(self, tp_size, ep_size, attention_dp,
+                          overlap_scheduler, use_msa, cuda_graph):
+        if use_msa:
+            from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import \
+                msa_package_available
+            if not msa_package_available():
+                pytest.skip("MSA kernels (fmha_sm100) not available")
+        model_name = "nvidia/MiniMax-M3-NVFP4"
+        model_path = f"{llm_models_root()}/MiniMax-M3-NVFP4"
+        max_draft_len = 3
+        spec_config = Eagle3DecodingConfig(
+            max_draft_len=max_draft_len,
+            speculative_model=f"{llm_models_root()}/MiniMax-M3-EAGLE3",
+        )
+        # The runtime forces tokens_per_block per implementation (128 MSA / 32
+        # reference).
+        kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.6,
+                                        enable_block_reuse=False)
+        with LLM(
+                model_path,
+                tensor_parallel_size=tp_size,
+                moe_expert_parallel_size=ep_size,
+                kv_cache_config=kv_cache_config,
+                sparse_attention_config=MiniMaxM3SparseAttentionConfig(
+                    implementation="msa" if use_msa else "triton"),
+                moe_config=MoeConfig(backend="CUTLASS"),
+                max_seq_len=4096,
+                # The fmha_sm100 decode planner caps total_q x num_qo_heads at
+                # 65536; with 1 + draft_len = 4 verify tokens per row that
+                # bounds the batch at 1024 / TP-sharded heads (256 unsharded
+                # under attention DP).
+                max_batch_size=256 if attention_dp else 512,
+                speculative_config=spec_config,
+                # Graphs + spec requires the MSA path: its verify batches
+                # are decode-shaped and capture-safe (the reference path
+                # rejects graphs+spec at creation).
+                cuda_graph_config=CudaGraphConfig(
+                    enable_padding=True,
+                    max_batch_size=64 if attention_dp else 128,
+                ) if cuda_graph else None,
+                disable_overlap_scheduler=not overlap_scheduler,
+                enable_attention_dp=attention_dp,
+                enable_iter_perf_stats=True,
+                trust_remote_code=True) as llm:
+            assert llm.args.quant_config.quant_algo == QuantAlgo.MIXED_PRECISION
+
+            def drain_spec_stats(llm):
+                drafted = accepted = steps = 0
+                for s in llm.get_stats(timeout=2):
+                    s = json.loads(s) if isinstance(s, str) else s
+                    sd = s.get("specDecodingStats") or {}
+                    drafted += sd.get("numDraftTokens", 0)
+                    accepted += sd.get("numAcceptedTokens", 0)
+                    steps += sd.get("numRequestsWithDraftTokens", 0)
+                return drafted, accepted, steps
+
+            task = MMLU(model_name)
+            task.evaluate(llm)
+            task = GSM8K(model_name)
+            task.evaluate(llm)
+
+            # Chat-format acceptance — the drafter's training distribution
+            # (Inferact/MiniMax-M3-EAGLE3 card: 0.839 / 3.518). Reuses the
+            # live engine and the cached dataset; ~20 s under CUDA graphs.
+            questions = [
+                r["question"]
+                for r in load_dataset("gsm8k", "main", split="test")
+            ][:200]
+            chat_prompts = [
+                llm.tokenizer.apply_chat_template([{
+                    "role": "user",
+                    "content": q
+                }],
+                                                  tokenize=False,
+                                                  add_generation_prompt=True)
+                for q in questions
+            ]
+            drain_spec_stats(llm)
+            llm.generate(chat_prompts,
+                         SamplingParams(max_tokens=512, temperature=0))
+            drafted, accepted, steps = drain_spec_stats(llm)
+            assert steps > 0, "no speculative iterations recorded"
+            chat_rate = accepted / drafted
+            chat_length = 1 + accepted / steps
+            print(f"MiniMax-M3 Eagle3 chat-GSM8K acceptance: rate="
+                  f"{chat_rate:.3f}, mean acceptance length="
+                  f"{chat_length:.3f} ({steps} spec iterations)")
+            assert chat_rate > 0.78, \
+                f"Eagle3 chat-GSM8K acceptance rate too low: {chat_rate:.3f}"
+            assert chat_length > 3.3, \
+                f"Eagle3 chat-GSM8K acceptance length too low: {chat_length:.3f}"
+
 
 @skip_pre_blackwell
 class TestGLM5FP8(LlmapiAccuracyTestHarness):
diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt
index fa1882623a3a..684645567645 100644
--- a/tests/integration/test_lists/qa/llm_function_core.txt
+++ b/tests/integration/test_lists/qa/llm_function_core.txt
@@ -665,6 +665,8 @@ accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[use_msa=False] TIMEO
 accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph[use_msa=False] TIMEOUT (180)
 accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=False] TIMEOUT (180)
 accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=True] TIMEOUT (60)
+accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True-use_msa=True-cuda_graph=True] TIMEOUT (180)
+accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=True-overlap_scheduler=True-use_msa=True-cuda_graph=True] TIMEOUT (180)
 accuracy/test_llm_api_pytorch.py::TestMinistral8BInstruct::test_auto_dtype
 accuracy/test_llm_api_pytorch.py::TestMinistral8BInstruct::test_fp8
 accuracy/test_llm_api_pytorch.py::TestMistralLarge3_675B::test_fp8[latency_moe_deepgemm]
diff --git a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py
index 58d651631608..07c8c8abe321 100644
--- a/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py
+++ b/tests/unittest/_torch/attention/sparse/test_minimax_m3_msa_backend.py
@@ -74,3 +74,27 @@ def test_msa_proxy_max_score_view_is_contiguous_over_stable_store():
     # Oversized requests are rejected rather than silently corrupting memory.
     with pytest.raises(ValueError, match=r"msa_max_score backing store"):
         metadata.msa_proxy_max_score_view(num_index_heads, worst_k, max_batch + 1)
+
+
+def test_per_token_valid_blocks_multi_token_decode():
+    """Spec-verify decode rows expose one entry per query TOKEN, walking the
+    causal ladder within the verify window."""
+    from tensorrt_llm._torch.attention_backend.sparse.minimax_m3.msa_utils import (
+        per_token_valid_blocks,
+    )
+
+    # One request verifying 4 tokens against kv_len 10 (offset 6): token t
+    # attends 7 + t positions; with 2-token blocks that is ceil((7+t)/2).
+    qo = torch.tensor([4], dtype=torch.int32)
+    kv = torch.tensor([10], dtype=torch.int32)
+    off = torch.tensor([6], dtype=torch.int32)
+    n_valid = per_token_valid_blocks(qo, kv, off, causal=True, block_size=2)
+    assert n_valid.tolist() == [4, 4, 5, 5]
+
+    # Mixed batch: an ordinary decode row (qo=1) alongside a verify row.
+    qo = torch.tensor([1, 3], dtype=torch.int32)
+    kv = torch.tensor([9, 6], dtype=torch.int32)
+    off = kv - qo
+    n_valid = per_token_valid_blocks(qo, kv, off, causal=True, block_size=4)
+    # Row 0: 9 positions -> 3 blocks. Row 1 tokens attend 4, 5, 6 -> 1, 2, 2.
+    assert n_valid.tolist() == [3, 1, 2, 2]

From e36496b1a47d1708101f7794ac51c84410f228ee Mon Sep 17 00:00:00 2001
From: Zheyu Fu <zheyuf@NVIDIA.com>
Date: Thu, 23 Jul 2026 20:40:15 +0000
Subject: [PATCH 2/2] [TRTLLM-14093][fix] Thread spec_metadata through
 MTPEagleDynamicTreeWorker spec-dec prep

The Eagle3 one-model change in this PR widened
Eagle3OneModelWorker._prepare_attn_metadata_for_spec_dec to take an extra
spec_metadata argument (for CUDA-graph-safe kv_lens_cuda save/restore).
MTPEagleDynamicTreeWorker (via MTPEagleWorker) inherits that base but keeps
its own single-arg override whose body calls super() with only attn_metadata,
so the dynamic-tree MTP path raises TypeError: missing 1 required positional
argument 'spec_metadata' at runtime. Thread spec_metadata through the override,
its super() call, and the call site to match the base signature.

Signed-off-by: Zheyu Fu <zheyuf@NVIDIA.com>
---
 tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py
index 053d23362f98..e0f933f73217 100644
--- a/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py
+++ b/tensorrt_llm/_torch/speculative/mtp_dynamic_tree.py
@@ -169,8 +169,8 @@ def __init__(
         sm = get_sm_version()
         self._needs_mask_repack = sm < 100 or sm in (120, 121)
 
-    def _prepare_attn_metadata_for_spec_dec(self, attn_metadata):
-        super()._prepare_attn_metadata_for_spec_dec(attn_metadata)
+    def _prepare_attn_metadata_for_spec_dec(self, attn_metadata, spec_metadata):
+        super()._prepare_attn_metadata_for_spec_dec(attn_metadata, spec_metadata)
 
         batch_size = attn_metadata.num_seqs
         if hasattr(attn_metadata, "kv_lens_cuda"):
@@ -684,7 +684,7 @@ def _forward_impl(
         # Save attn/spec metadata before the draft loop mutates it.
         original_all_rank_num_tokens = attn_metadata.all_rank_num_tokens
         original_force_prepare_spec_dec_tree_mask = attn_metadata.force_prepare_spec_dec_tree_mask
-        self._prepare_attn_metadata_for_spec_dec(attn_metadata)
+        self._prepare_attn_metadata_for_spec_dec(attn_metadata, spec_metadata)
         attn_metadata.force_prepare_spec_dec_tree_mask = True
 
         # (c) Run the MTP draft tree loop -> build + store the next tree.