From 2c2e636373232255be458ec065672948a75cb02d Mon Sep 17 00:00:00 2001
From: Zheyu Fu <zheyuf@nvidia.com>
Date: Mon, 7 Sep 2026 16:16:45 -0700
Subject: [PATCH 1/7] [TRTLLM-14093][feat] One-model Eagle3 speculative
 decoding for MiniMax-M3 on the MSA backend

Add one-model Eagle3 support to the MiniMax-M3 MSA sparse attention path:
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), and a sync-free
on_update_kv_lens that re-derives slots, counts and the plans' length mirrors
on device after the overlap scheduler corrects kv_lens, clamped to the staged
lengths (which 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 steps are unchanged. The dense SDPA
reference path gains the causal-ladder verify mask.

Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
---
 docs/source/models/supported-models.md        |   2 +-
 .../backends/sparse/minimax_m3/msa_backend.py | 288 ++++++++++++++++--
 .../sparse/minimax_m3/triton_metadata.py      | 174 +++++++----
 .../_torch/models/modeling_minimaxm3.py       | 117 ++++---
 .../defs/accuracy/references/gsm8k.yaml       |   4 +
 .../defs/accuracy/references/mmlu.yaml        |   4 +
 .../attention/sparse/msa/test_msa_backend.py  | 212 ++++++++++++-
 .../hf/test_minimaxm3_weight_mapper.py        |   2 +
 8 files changed, 657 insertions(+), 146 deletions(-)

diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md
index 4071dfeef728..e6b2dc9d8f7b 100644
--- a/docs/source/models/supported-models.md
+++ b/docs/source/models/supported-models.md
@@ -91,7 +91,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl
 | `Gemma4ForConditionalGeneration` | Untested          | Yes        | Untested                   | No                    | Yes             | MTP                        | Yes           | Untested         | No             | Yes                      | Untested              | Untested        |
 | `Gemma4UnifiedForConditionalGeneration` | Untested          | Untested   | Untested                   | No                    | Yes             | No                         | Yes           | Untested         | No             | Yes                      | Untested              | Untested        |
 | `Step3p7ForConditionalGeneration`| Yes               | Yes        | Yes                        | Untested              | Untested        | MTP                        | Yes           | Untested         | Untested       | Yes                      | Untested              | Untested        |
-| `MiniMaxM3SparseForConditionalGeneration` [^12] | Yes               | Yes        | Yes                        | Untested              | Untested        | No                         | Yes           | Untested         | No             | N/A                      | Untested              | Untested        |
+| `MiniMaxM3SparseForConditionalGeneration` [^12] | Yes               | Yes        | Yes                        | Untested              | Untested        | EAGLE-3 (Linear)           | 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.
diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py
index 93b0d9af8f5d..c41b4d7c17cb 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py
@@ -16,6 +16,9 @@
     into them. The standard CUDAGraphRunner clones one metadata per graph
     batch size (create_cuda_graph_metadata), so no per-batch-size cache is
     needed here.
+  * With Eagle3 a decode row has 1 + draft_len query tokens. Slots, valid-block
+    counts and plan rows are per token; on_update_kv_lens re-derives them after
+    the overlap scheduler corrects kv_lens on device.
 
 The classes subclass TrtllmAttention and TrtllmAttentionMetadata, imported at
 module scope. This is cycle-free because the fmha registry defers its
@@ -115,6 +118,15 @@ def _worst_case_proxy_max_k_tiles(
     "workspace_o",
     "workspace_lse",
 )
+# Per-row length fields on_update_kv_lens patches in a plan. Dense plans use
+# kv_segment_lens/qo_offset; sparse plans (tagged "MM-SA-Nv") use seqused_k.
+_MSA_DENSE_LENGTH_KEYS = ("kv_segment_lens", "qo_offset")
+_MSA_SPARSE_LENGTH_KEYS = ("seqused_k",)
+
+
+def _msa_plan_length_keys(sub_plan: dict) -> tuple:
+    """Length fields to patch in one fmha_sm100 sub-plan."""
+    return _MSA_SPARSE_LENGTH_KEYS if sub_plan.get("MM-SA-Nv") else _MSA_DENSE_LENGTH_KEYS
 
 
 class _MsaGraphSafePlan:
@@ -223,6 +235,10 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata):
     need no graph-stable storage. Plans are built in _build_step_plans:
     pure-decode batches use the graph-safe owners (msa_decode_*_plan) while
     prefill/mixed batches keep plain eager tuples (msa_eager_*_plan).
+
+    With Eagle3 a decode row has 1 + draft_len query tokens: slots, valid-block
+    counts and plan rows are per token, and the overlap scheduler corrects
+    kv_lens on device after prepare(); see on_update_kv_lens.
     """
 
     # Graph-stable buffers; consumers slice to the live count at the call
@@ -231,6 +247,13 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata):
     msa_kv_indices: Optional[torch.Tensor] = None
     msa_max_score: Optional[torch.Tensor] = None
     msa_n_valid_blocks: Optional[torch.Tensor] = None
+    # Inputs for on_update_kv_lens: the step's request->slot table, each query
+    # token's request row and offset within the request, and the kv_lens
+    # prepare() staged (the upper bound the correction is clamped to).
+    msa_req_to_token: Optional[torch.Tensor] = None
+    msa_q_batch_row: Optional[torch.Tensor] = None
+    msa_q_intra: Optional[torch.Tensor] = None
+    msa_kv_lens_staged: Optional[torch.Tensor] = None
 
     # _msa_buffers_ready gates the once-only device buffers;
     # _msa_fields_ready marks that the current step's buffers are populated.
@@ -257,6 +280,12 @@ class MiniMaxM3MsaSparseAttentionMetadata(TrtllmAttentionMetadata):
     # persistent backing store for the view.
     _msa_eager_n_valid_buf: Optional[torch.Tensor] = None
     _msa_eager_n_valid_blocks: Optional[torch.Tensor] = None
+    # Host-side token offset of each request (plus the total), so
+    # on_update_kv_lens can slice a sub-plan's token range without a device read.
+    _msa_q_token_starts: tuple = (0,)
+    # True only with speculative decoding; otherwise on_update_kv_lens and its
+    # staging are skipped and non-speculative steps run as before.
+    _msa_kv_lens_dynamic: bool = False
 
     def __post_init__(self) -> None:
         super().__post_init__()
@@ -282,11 +311,19 @@ 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).
+
+        Excludes num_extra_kv_tokens (speculative draft slots): MSA needs the
+        attended length.
+        """
         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
         if out.dtype != torch.int32:
             out = out.to(torch.int32)
         return maybe_pin_memory(out)
@@ -382,6 +419,36 @@ def _create_msa_buffers(self) -> None:
             dtype=torch.int32,
             capture_graph=capture_graph,
         )
+        # Inputs for on_update_kv_lens.
+        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_kv_lens_staged = self.get_empty(
+            buffers,
+            (max_num_sequences,),
+            cache_name="msa_kv_lens_staged",
+            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
@@ -397,37 +464,49 @@ 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 query tokens in one decode step (1 + draft_len per row).
+
+        Capped at 16384 so total_q * num_qo_heads stays under fmha_sm100's
+        65536 planner limit with 4 index heads.
+        """
+        max_seqs = int(self.max_num_sequences)
+        max_toks = int(self.max_num_tokens or 0)
+        if max_toks <= 0:
+            return max_seqs
+        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.
+        Sized for the worst-case max_k_tiles and query-token count (more than
+        the batch size under speculative verify), so one allocation serves every
+        decode step. msa_proxy_max_score_view slices the per-step shape.
         """
         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,
@@ -442,14 +521,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
 
@@ -471,7 +551,7 @@ 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,
         )
@@ -495,6 +575,126 @@ def prepare(self) -> None:
         self._build_msa_fields()
         self._build_step_plans()
 
+    def _msa_kv_lens_may_change(self) -> bool:
+        """Whether kv_lens can change after prepare(): only with speculative decoding.
+
+        max_total_draft_tokens is set by the engine's update_spec_dec_param; the
+        other checks mirror TrtllmAttentionMetadata's spec_active.
+        """
+        params = self.kv_cache_params
+        runtime_features = self.runtime_features
+        return bool(
+            self.max_total_draft_tokens
+            or self.draft_kv_cache_manager is not None
+            or self.is_spec_decoding_enabled
+            or (params is not None and params.num_extra_kv_tokens)
+            or (runtime_features is not None and runtime_features.has_speculative_draft_tokens)
+        )
+
+    def _msa_live_plans(self) -> tuple:
+        """Plans live this step: the graph-safe owners on decode, else the eager tuples."""
+        live = []
+        for owner, eager in (
+            (self._msa_proxy_plan, self._msa_eager_proxy_plan),
+            (self._msa_gqa_plan, self._msa_eager_gqa_plan),
+            (self._msa_dense_plan, self._msa_eager_dense_plan),
+        ):
+            plan = owner.plan if owner is not None else None
+            if plan is None:
+                plan = eager
+            if plan is not None:
+                live.append(plan)
+        return tuple(live)
+
+    def on_update_kv_lens(self) -> None:
+        """Re-derive slots, valid-block counts and plan lengths from the corrected kv_lens_cuda.
+
+        The overlap scheduler shortens kv_lens on device after prepare() staged
+        full-acceptance values (the draft loop does the same between CUDA-graph
+        warmup runs). Shrinking keeps the staged page table and worklists
+        valid, so only per-row lengths are patched; the clamp to
+        msa_kv_lens_staged enforces that. Device-only, capture-safe and
+        idempotent; skipped without speculative decoding.
+        """
+        super().on_update_kv_lens()
+        if not self._msa_fields_ready or not self._msa_kv_lens_dynamic:
+            return
+        batch = int(self.num_seqs)
+        total_q = int(self.num_tokens)
+        if batch <= 0 or total_q <= 0:
+            return
+        kv_true = torch.minimum(self.kv_lens_cuda[:batch], self.msa_kv_lens_staged[:batch])
+        qbr = self.msa_q_batch_row[:total_q].to(torch.long)
+        qo_dev = self.seq_lens_cuda[:batch]
+        kv_true_tok = kv_true[qbr]
+        # Position each query token attends up to.
+        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. clamp_min(1) keeps padding
+        # rows from masking every block, which would NaN the GQA row.
+        page = int(self.kv_cache_manager.tokens_per_block)
+        n_valid = torch.div((pos + 1).clamp_min(1) + (page - 1), page, rounding_mode="floor")
+        n_valid_buf = (
+            self._msa_eager_n_valid_blocks
+            if self._msa_eager_n_valid_blocks is not None
+            else self.msa_n_valid_blocks
+        )
+        if n_valid_buf is not None:
+            n_valid_buf[:total_q].copy_(n_valid.to(torch.int32))
+
+        # Patch each sub-plan's length rows. A plan is (has_mixed, split,
+        # batch, decode_sub, prefill_sub); the sub-plans cover rows [0, split)
+        # and [split, batch) and the names are positional only. A sub-plan has
+        # one row per request or one per query token, so the row count picks
+        # the source. qo_offset must stay >= 0 (negative hits a kernel sentinel).
+        per_request = {
+            "kv_segment_lens": kv_true,
+            "qo_offset": (kv_true - qo_dev).clamp_min(0),
+            "seqused_k": kv_true,
+        }
+        per_token = {
+            "kv_segment_lens": kv_true_tok,
+            "qo_offset": pos.clamp_min(0),
+            "seqused_k": (pos + 1).clamp_min(0),
+        }
+        starts = self._msa_q_token_starts
+        for plan in self._msa_live_plans():
+            has_mixed, split, _, decode_sub, prefill_sub = plan
+            subs = (
+                ((decode_sub, 0, split), (prefill_sub, split, batch))
+                if has_mixed
+                else ((decode_sub, 0, batch),)
+            )
+            for sub, first, last in subs:
+                if sub is None:
+                    continue
+                tok_first, tok_last = starts[first], starts[last]
+                for key in _msa_plan_length_keys(sub):
+                    dst = sub.get(key)
+                    if dst is None:
+                        raise RuntimeError(
+                            f"MSA plan has no length mirror {key!r}, so the corrected "
+                            "kv_lens cannot reach the kernel."
+                        )
+                    rows = int(dst.shape[0])
+                    if rows == last - first:
+                        src = per_request[key][first:last]
+                    elif rows == tok_last - tok_first:
+                        src = per_token[key][tok_first:tok_last]
+                    else:
+                        raise RuntimeError(
+                            f"MSA plan {key!r} has {rows} rows for requests "
+                            f"[{first}, {last}); expected {last - first} (one per "
+                            f"request) or {tok_last - tok_first} (one per query token)."
+                        )
+                    dst.copy_(src)
+
     def _build_step_plans(self) -> None:
         """Build the three layer-invariant fmha_sm100 plans once per step.
 
@@ -536,7 +736,6 @@ def _build_step_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
@@ -596,14 +795,17 @@ def _build_step_plans(self) -> None:
             self._msa_eager_gqa_plan = gqa_plan
             self._msa_eager_dense_plan = dense_plan
             # Stage the valid-block count to the device once for the whole step
-            # (see _msa_eager_n_valid_blocks).
+            # (see _msa_eager_n_valid_blocks). clamp_min(1) as in
+            # on_update_kv_lens: a zero-valid row would NaN the GQA row.
             n_valid_host = per_token_valid_blocks(
                 qo_lens_cpu, kv_lens_cpu, qo_offset_cpu, causal=True, block_size=page_size
             )
             total_q = int(n_valid_host.shape[0])
             if total_q > 0:
                 dev_buf = self._ensure_eager_n_valid_buffer(total_q, device)
-                dev_buf[:total_q].copy_(n_valid_host.to(torch.int32), non_blocking=True)
+                dev_buf[:total_q].copy_(
+                    n_valid_host.clamp_min(1).to(torch.int32), non_blocking=True
+                )
                 self._msa_eager_n_valid_blocks = dev_buf[:total_q]
             return
 
@@ -616,27 +818,29 @@ def _build_step_plans(self) -> None:
         )
 
         # Allocate the graph-safe plan owners once per metadata; later steps
-        # only refresh their contents below.
+        # only refresh their contents below. Worklists have one row per query
+        # token under speculative verify, so size them by tokens, not batch.
         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,
             )
@@ -650,14 +854,17 @@ def _build_step_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.
+        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.
 
         The page table and per-new-token cache slots are derived via the
         build_paged_kv_slot_mapping helper, then copied into the persistent
-        buffers. The transient builder tensors are discarded.
+        buffers. The transient builder tensors are discarded. With speculative
+        decoding the inputs for on_update_kv_lens are staged as well.
         """
         self._msa_fields_ready = False
         if not self._msa_buffers_ready:
@@ -676,18 +883,10 @@ 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
-        # (num_cached) and decode (kv_len - 1 with qo_len 1).
+        # (num_cached) and decode (kv_len - qo_len).
         mapping = build_paged_kv_slot_mapping(
             kv_cache_manager=kv_cache_manager,
             request_ids=request_ids,
@@ -695,6 +894,7 @@ def _build_msa_fields(self) -> None:
             qo_offset_cpu=qo_offset_cpu,
             device=cache_device,
         )
+        req_to_token = mapping.req_to_token
         out_cache_loc = mapping.out_cache_loc
         # The page table comes from the same host block ids the mapping was
         # built from, so it costs no device work.
@@ -715,6 +915,36 @@ 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)
+
+        self._msa_kv_lens_dynamic = self._msa_kv_lens_may_change()
+        if not self._msa_kv_lens_dynamic:
+            self._msa_fields_ready = True
+            return
+
+        # Inputs for on_update_kv_lens: the request -> slot table and each query
+        # token's (request row, offset in request). Pinned and non-blocking so
+        # the copies do not synchronize the stream.
+        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_(
+            maybe_pin_memory(batch_row_cpu), non_blocking=True
+        )
+        self.msa_q_intra[:total_new_tokens].copy_(maybe_pin_memory(intra_cpu), non_blocking=True)
+        # The staged lens are the upper bound on_update_kv_lens clamps to.
+        self.msa_kv_lens_staged[:batch_size].copy_(
+            self.kv_lens_cuda[:batch_size], non_blocking=True
+        )
+        # Token offset of each request, plus the total.
+        self._msa_q_token_starts = (0, *torch.cumsum(qo_long, 0).tolist())
         self._msa_fields_ready = True
 
     def msa_idx_k_cache(self, layer_idx: int) -> torch.Tensor:
diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_metadata.py
index 696f0f68055c..0390e5a158b8 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_metadata.py
+++ b/tensorrt_llm/_torch/attention/backends/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,8 @@ 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 in decode: 1, or 1 + draft_len with Eagle3.
+    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 +167,10 @@ 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())
+            # With the overlap scheduler, optimistic lengths can run past the page
+            # table; SDPA uses max_seqlen_k as the mask width, so clamp it.
+            self.max_seqlen_k = min(max_k, int(self.req_to_token.shape[1]))
 
 
 def ensure_metadata_on_device(
@@ -369,6 +374,40 @@ 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-query-token K positions and KV slot ids, computed on device.
+
+    Shared by the metadata builder and on_update_kv_lens so they agree.
+    """
+    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 may point past the last allocated page. Those
+    # slots are placeholders that on_update_kv_lens fixes before use, but the
+    # gather must stay in bounds. Not in-place: .to() may alias the input.
+    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 (the new token sits at ``seq_lens[b] - 1``).
+
+    Same clamp as derive_q_positions_and_cache_slots; ``min=0`` also covers
+    empty dummy rows.
+    """
+    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 +578,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 +594,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 +613,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 +643,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 +651,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 +666,12 @@ 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` (like
+    ``DSAtrtllmAttentionMetadata``) so one-model Eagle3 draft layers can run
+    :class:`TrtllmAttention` on this metadata.
 
     Overrides :meth:`prepare` so the M3-sparse
     :class:`MiniMaxM3TritonSparseAttentionMetadata` and the per-new-token
@@ -808,23 +827,9 @@ def prepare(self) -> None:
 
         static_buffers = self._maybe_get_m3_static_buffers(cache_device, kv_cache_manager)
 
-        # Any batch containing a context (prefill or chunked extend)
-        # request takes the extend path. For prefill rows
-        # ``num_cached_per_seq`` is ``prefix_lens`` and the full new
-        # chunk is ``extend_seq_len``; for decode rows
-        # ``num_cached`` is ``kv_len - 1`` and ``extend_seq_len`` is
-        # 1, so the same builder produces the correct one-slot
-        # entry. Pure-decode batches (``num_contexts == 0``) still
-        # take the decode optimization for CUDA-graph warmup
-        # geometry.
-        #
-        # Mixed prefill+decode batches always take the extend path:
-        # the prefill kernel handles decode rows as 1-slot extends.
-        # The decode branch below is a pure-decode-only perf
-        # 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 generation rows (spec verify) also take the extend path;
+        # plain 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 +867,50 @@ 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``.
+
+        With the overlap scheduler and speculative decoding, prepare() runs
+        with optimistic lengths and the engine corrects ``kv_lens_cuda`` on
+        device before calling this (as DSAtrtllmAttentionMetadata does).
+        Device-only and idempotent. ``seq_lens_cpu`` / ``max_seqlen_k`` keep
+        the optimistic values; they only bound widths 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 layout
+            # (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:
+            # No-op today (decode rows are not corrected); kept for symmetry.
+            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 1176262d960e..731ef7dbdff3 100644
--- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py
+++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py
@@ -64,16 +64,12 @@
 from ..modules.rms_norm import RMSNorm
 from ..moe.fused_moe import MiniMaxM3MoeRoutingMethod, SwigluBiasActivation, create_moe
 from ..pyexecutor.breakable_cuda_graph import eager_on_graph, is_in_breakable_cuda_graph
+from ..speculative import SpecMetadata
 from ..utils import AuxStreamType, EventType, get_model_extra_attrs, is_torch_compiling
 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,
-    filter_weights,
-    register_auto_model,
-)
+from .modeling_speculative import SpecDecOneEngineForCausalLM
+from .modeling_utils import DecoderModel, ModelConfig, filter_weights, 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,
@@ -1211,7 +1207,14 @@ def _sdpa_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)
+        # Under CUDA graphs bake a fixed gather/mask width: min(page-table width,
+        # max_seq_len). The raw table width is inflated by the KV-estimation pass
+        # and would OOM the gather; the seq_lens mask hides 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
@@ -1238,15 +1241,11 @@ def _sdpa_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
@@ -1255,6 +1254,9 @@ def _sdpa_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
@@ -1295,35 +1297,37 @@ def _sdpa_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 row b attends
+            # seq_lens[b] - qo_len + t + 1 positions (qo_len=1 is the usual mask).
+            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: all heads at once needs an
+            # O(batch * max_k * num_heads) temporary that overflows the CUDA-graph
+            # pool under attention DP.
+            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; transpose(1, 2)
+            # .reshape would scramble (head, head_dim) when H != head_dim.
+            output.view(batch, qo_len, self.num_heads, self.head_dim).copy_(out_b.transpose(1, 2))
 
         return output
 
@@ -1747,6 +1751,7 @@ def forward(
         hidden_states: torch.Tensor,
         attn_metadata: AttentionMetadata,
         residual: Optional[torch.Tensor],
+        spec_metadata: Optional[SpecMetadata] = None,
         **kwargs,
     ) -> torch.Tensor:
         # Layer-0 prologue only. For every subsequent layer the input_layernorm
@@ -1772,10 +1777,19 @@ def forward(
             **kwargs,
         )
 
+        if spec_metadata is not None and spec_metadata.is_layer_capture(self.layer_idx):
+            # Eagle3 captures this layer's raw output, so keep the feed-forward
+            # AllReduce in the module and the boundary norm unfused (as DeepSeek-V3
+            # does on capture layers).
+            self.post_feed_forward_fusion = False
         if self.block_sparse_moe is not None:
-            hidden_states, residual = self.forward_MoE(hidden_states, attn_metadata, residual)
+            hidden_states, residual = self.forward_MoE(
+                hidden_states, attn_metadata, residual, spec_metadata=spec_metadata
+            )
         else:
-            hidden_states, residual = self.forward_mlp(hidden_states, residual)
+            hidden_states, residual = self.forward_mlp(
+                hidden_states, residual, spec_metadata=spec_metadata
+            )
 
         return hidden_states, residual
 
@@ -1849,6 +1863,7 @@ def forward_MoE(
         hidden_states: torch.Tensor,
         attn_metadata: AttentionMetadata,
         residual: torch.Tensor,
+        spec_metadata: Optional[SpecMetadata] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         hidden_states, residual = self._apply_pre_feed_forward_norm(hidden_states, residual)
 
@@ -1858,6 +1873,8 @@ def forward_MoE(
             final_all_reduce_params=self._feed_forward_all_reduce_params(),
         )
 
+        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)
         hidden_states, residual = self._apply_next_layer_layernorm(hidden_states, residual)
         return hidden_states, residual
 
@@ -1865,6 +1882,7 @@ def forward_mlp(
         self,
         hidden_states: torch.Tensor,
         residual: torch.Tensor,
+        spec_metadata: Optional[SpecMetadata] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         hidden_states, residual = self._apply_pre_feed_forward_norm(hidden_states, residual)
 
@@ -1873,6 +1891,8 @@ def forward_mlp(
             final_all_reduce_params=self._feed_forward_all_reduce_params(),
         )
 
+        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)
         hidden_states, residual = self._apply_next_layer_layernorm(hidden_states, residual)
         return hidden_states, residual
 
@@ -1936,6 +1956,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):
@@ -1952,6 +1973,7 @@ def forward(
                 hidden_states=hidden_states,
                 attn_metadata=attn_metadata,
                 residual=residual,
+                spec_metadata=spec_metadata,
             )
 
         # When setup_aliases has chained the final norm into the last decoder
@@ -2031,7 +2053,7 @@ def _fold_gemma_boundary_norm_weights(weights):
 
 
 @register_auto_model("MiniMaxM3SparseForCausalLM")
-class MiniMaxM3ForCausalLM(DecoderModelForCausalLM[MiniMaxM3Model, PretrainedConfig]):
+class MiniMaxM3ForCausalLM(SpecDecOneEngineForCausalLM[MiniMaxM3Model, PretrainedConfig]):
     """Text-only M3 model."""
 
     @classmethod
@@ -2059,12 +2081,7 @@ 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/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml
index b77f76c72638..09527c19d1f0 100644
--- a/tests/integration/defs/accuracy/references/gsm8k.yaml
+++ b/tests/integration/defs/accuracy/references/gsm8k.yaml
@@ -521,6 +521,10 @@ nvidia/MiniMax-M3-NVFP4:
   - quant_algo: MIXED_PRECISION
     kv_cache_quant_algo: FP8
     accuracy: 86
+  - quant_algo: MIXED_PRECISION
+    kv_cache_quant_algo: FP8
+    spec_dec_algo: Eagle3
+    accuracy: 86
 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 9ee6b600f196..165ceba141d4 100644
--- a/tests/integration/defs/accuracy/references/mmlu.yaml
+++ b/tests/integration/defs/accuracy/references/mmlu.yaml
@@ -185,6 +185,10 @@ nvidia/MiniMax-M3-NVFP4:
   - quant_algo: MIXED_PRECISION
     kv_cache_quant_algo: FP8
     accuracy: 81
+  - quant_algo: MIXED_PRECISION
+    kv_cache_quant_algo: FP8
+    spec_dec_algo: Eagle3
+    accuracy: 81
 moonshotai/Kimi-K2-Instruct:
   - quant_algo: FP8_BLOCK_SCALES
     accuracy: 87.65
diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py
index 62884e61a626..fbbd9ca0a53f 100644
--- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py
+++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py
@@ -187,9 +187,12 @@ def test_msa_metadata_rejects_undersized_max_score_buffer():
     metadata_cls = MiniMaxM3MsaSparseAttention.Metadata
     metadata = metadata_cls.__new__(metadata_cls)
     # Flat backing store sized for 4 heads * 8 k-tiles * 2 batch = 64 elements,
-    # too small for the plan's required 4 * 16 * 2 = 128.
+    # too small for the plan's required 4 * 16 * 2 = 128. One query token per
+    # decode row, so the token bound equals the batch here.
     metadata.msa_max_score = torch.zeros(4 * 8 * 2)
     metadata.kv_cache_manager = None
+    metadata.max_num_sequences = 2
+    metadata.max_num_tokens = 2
 
     with pytest.raises(ValueError, match=r"msa_max_score backing store"):
         metadata._ensure_msa_decode_scratch_buffers(
@@ -299,6 +302,31 @@ def test_msa_paged_hnd_input_materializes_unaligned_outer_stride() -> None:
     torch.testing.assert_close(prepared, view)
 
 
+def test_per_token_valid_blocks_multi_token_decode():
+    """Spec-verify decode rows get one entry per query token, following the causal
+    ladder inside the verify window.
+    """
+    from tensorrt_llm._torch.attention.backends.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]
+
+
 def test_msa_index_k_uses_hnd_cache_view_and_writer():
     metadata_cls = MiniMaxM3MsaSparseAttention.Metadata
     metadata = metadata_cls.__new__(metadata_cls)
@@ -732,3 +760,185 @@ def get_block_ids_per_seq(self, request_ids):
     )
     for b, slot in enumerate(padded.out_cache_loc.tolist()):
         assert slot in req_to_token[b].tolist()
+
+
+def test_msa_scratch_sizing_covers_spec_verify_tokens():
+    """With Eagle3 a decode step has 1 + draft_len query tokens per request, so the
+    proxy scratch must be sized by tokens, not by batch.
+    """
+    metadata_cls = MiniMaxM3MsaSparseAttention.Metadata
+    metadata = metadata_cls.__new__(metadata_cls)
+    metadata.kv_cache_manager = None
+    # 2 sequences, 4 tokens each (draft_len=3): 8 decode tokens per step.
+    metadata.max_num_sequences = 2
+    metadata.max_num_tokens = 8
+    # Store sized for batch-only sizing (4 heads * 16 k-tiles * 2), which is
+    # too small once tokens are accounted for (4 * 16 * 8).
+    metadata.msa_max_score = torch.zeros(4 * 16 * 2)
+
+    with pytest.raises(ValueError, match=r"msa_max_score backing store"):
+        metadata._ensure_msa_decode_scratch_buffers(
+            num_index_heads=4,
+            max_batch=2,
+            capture_graph=False,
+            required_max_k_tiles=16,
+        )
+
+
+def _kv_lens_update_metadata(monkeypatch, *, qo_lens, kv_staged, kv_corrected, page_size):
+    """Metadata with just enough staged state for on_update_kv_lens: two requests
+    with several query tokens each, optimistic kv_lens from prepare(), then the
+    corrected kv_lens. Everything lives on the host.
+    """
+    from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import msa_backend
+
+    # The base hook only touches MLA caches this metadata never builds.
+    monkeypatch.setattr(msa_backend.TrtllmAttentionMetadata, "on_update_kv_lens", lambda self: None)
+    metadata_cls = MiniMaxM3MsaSparseAttention.Metadata
+    metadata = metadata_cls.__new__(metadata_cls)
+    qo = torch.tensor(qo_lens, dtype=torch.int32)
+    batch = int(qo.shape[0])
+    total_q = int(qo.sum())
+    metadata._seq_lens = qo
+    metadata._seq_lens_cuda = qo.clone()
+    metadata._num_tokens = total_q
+    metadata.kv_lens_cuda = torch.tensor(kv_corrected, dtype=torch.int32)
+    metadata.msa_kv_lens_staged = torch.tensor(kv_staged, dtype=torch.int32)
+    metadata.kv_cache_manager = SimpleNamespace(tokens_per_block=page_size)
+    # Slot table: request b holds slot 100 * b + position.
+    width = 16
+    metadata.msa_req_to_token = torch.arange(width, dtype=torch.int32).unsqueeze(
+        0
+    ) + 100 * torch.arange(batch, dtype=torch.int32).unsqueeze(1)
+    qo_long = qo.to(torch.long)
+    metadata.msa_q_batch_row = torch.repeat_interleave(
+        torch.arange(batch, dtype=torch.int32), qo_long
+    )
+    starts = torch.cumsum(qo_long, 0) - qo_long
+    metadata.msa_q_intra = (
+        torch.arange(total_q, dtype=torch.int64) - torch.repeat_interleave(starts, qo_long)
+    ).to(torch.int32)
+    metadata._msa_q_token_starts = (0, *torch.cumsum(qo_long, 0).tolist())
+    metadata.msa_out_cache_loc = torch.full((total_q + 3,), -1, dtype=torch.int32)
+    metadata.msa_n_valid_blocks = torch.zeros(total_q + 3, dtype=torch.int32)
+    metadata._msa_eager_n_valid_blocks = None
+    metadata._msa_proxy_plan = None
+    metadata._msa_gqa_plan = None
+    metadata._msa_dense_plan = None
+    metadata._msa_eager_proxy_plan = None
+    metadata._msa_eager_gqa_plan = None
+    metadata._msa_eager_dense_plan = None
+    metadata._msa_fields_ready = True
+    metadata._msa_kv_lens_dynamic = True
+    return metadata
+
+
+def test_on_update_kv_lens_rederives_slots_counts_and_plan_lengths(monkeypatch):
+    """Request 0 loses one rejected draft token (staged 9 -> corrected 8), request 1
+    is unchanged. Slots, valid-block counts and the plans' length rows must
+    follow, per request or per query token depending on the row count.
+    """
+    metadata = _kv_lens_update_metadata(
+        monkeypatch, qo_lens=(2, 3), kv_staged=(9, 12), kv_corrected=(8, 12), page_size=4
+    )
+    # Proxy plan: one row per request, dense length keys. Sparse GQA plan:
+    # row-expanded per query token, tagged "MM-SA-Nv" with seqused_k.
+    proxy_sub = {
+        "kv_segment_lens": torch.zeros(2, dtype=torch.int32),
+        "qo_offset": torch.zeros(2, dtype=torch.int32),
+    }
+    gqa_sub = {"MM-SA-Nv": True, "seqused_k": torch.zeros(5, dtype=torch.int32)}
+    metadata._msa_proxy_plan = SimpleNamespace(plan=(False, 2, 2, proxy_sub, None))
+    metadata._msa_gqa_plan = SimpleNamespace(plan=(False, 2, 2, gqa_sub, None))
+
+    metadata.on_update_kv_lens()
+
+    # Token positions: request 0 attends 8 tokens with 2 queries -> 6, 7;
+    # request 1 attends 12 with 3 queries -> 9, 10, 11.
+    assert metadata.msa_out_cache_loc[:5].tolist() == [6, 7, 109, 110, 111]
+    assert metadata.msa_out_cache_loc[5:].tolist() == [-1, -1, -1]
+    # ceil((pos + 1) / 4) per token.
+    assert metadata.msa_n_valid_blocks[:5].tolist() == [2, 2, 3, 3, 3]
+    assert proxy_sub["kv_segment_lens"].tolist() == [8, 12]
+    assert proxy_sub["qo_offset"].tolist() == [6, 9]
+    assert gqa_sub["seqused_k"].tolist() == [7, 8, 10, 11, 12]
+
+
+def test_on_update_kv_lens_clamps_to_the_staged_lens(monkeypatch):
+    """A correction can only shrink lengths; anything above the staged value is
+    clamped.
+    """
+    metadata = _kv_lens_update_metadata(
+        monkeypatch, qo_lens=(1, 1), kv_staged=(5, 7), kv_corrected=(9, 7), page_size=4
+    )
+    dense_sub = {
+        "kv_segment_lens": torch.zeros(2, dtype=torch.int32),
+        "qo_offset": torch.zeros(2, dtype=torch.int32),
+    }
+    metadata._msa_eager_dense_plan = (False, 2, 2, dense_sub, None)
+
+    metadata.on_update_kv_lens()
+
+    assert dense_sub["kv_segment_lens"].tolist() == [5, 7]
+    assert dense_sub["qo_offset"].tolist() == [4, 6]
+    assert metadata.msa_out_cache_loc[:2].tolist() == [4, 106]
+
+
+def test_on_update_kv_lens_patches_mixed_plans_per_sub_plan(monkeypatch):
+    """A mixed plan has two sub-plans over rows [0, split) and [split, batch); each is
+    patched over its own rows, per request or per query token.
+    """
+    metadata = _kv_lens_update_metadata(
+        monkeypatch, qo_lens=(1, 3), kv_staged=(6, 12), kv_corrected=(6, 11), page_size=4
+    )
+    decode_sub = {
+        "kv_segment_lens": torch.zeros(1, dtype=torch.int32),
+        "qo_offset": torch.zeros(1, dtype=torch.int32),
+    }
+    prefill_sub = {
+        "kv_segment_lens": torch.zeros(3, dtype=torch.int32),
+        "qo_offset": torch.zeros(3, dtype=torch.int32),
+    }
+    metadata._msa_eager_gqa_plan = (True, 1, 2, decode_sub, prefill_sub)
+
+    metadata.on_update_kv_lens()
+
+    assert decode_sub["kv_segment_lens"].tolist() == [6]
+    assert decode_sub["qo_offset"].tolist() == [5]
+    # Request 1 attends 11 tokens with 3 queries: positions 8, 9, 10.
+    assert prefill_sub["kv_segment_lens"].tolist() == [11, 11, 11]
+    assert prefill_sub["qo_offset"].tolist() == [8, 9, 10]
+
+
+def test_on_update_kv_lens_is_a_noop_without_speculative_decoding(monkeypatch):
+    """Without speculative decoding the patch and its staging are skipped. The gate
+    also sees max_total_draft_tokens, set by update_spec_dec_param, which
+    covers draft_len=1 under trtllm-gen (no extra KV tokens, linear-tree spec
+    decoding reported disabled).
+    """
+    metadata = _kv_lens_update_metadata(
+        monkeypatch, qo_lens=(1, 1), kv_staged=(4, 6), kv_corrected=(3, 5), page_size=4
+    )
+    metadata.max_total_draft_tokens = None
+    metadata.draft_kv_cache_manager = None
+    metadata.is_spec_decoding_enabled = False
+    metadata.kv_cache_params = SimpleNamespace(num_extra_kv_tokens=0)
+    metadata.runtime_features = None
+    assert metadata._msa_kv_lens_may_change() is False
+    metadata.max_total_draft_tokens = 1
+    assert metadata._msa_kv_lens_may_change() is True
+    metadata.max_total_draft_tokens = None
+    metadata.kv_cache_params = SimpleNamespace(num_extra_kv_tokens=2)
+    assert metadata._msa_kv_lens_may_change() is True
+
+    metadata._msa_kv_lens_dynamic = False
+    dense_sub = {
+        "kv_segment_lens": torch.zeros(2, dtype=torch.int32),
+        "qo_offset": torch.zeros(2, dtype=torch.int32),
+    }
+    metadata._msa_eager_dense_plan = (False, 2, 2, dense_sub, None)
+
+    metadata.on_update_kv_lens()
+
+    assert metadata.msa_out_cache_loc.tolist() == [-1] * 5
+    assert dense_sub["kv_segment_lens"].tolist() == [0, 0]
diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py b/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py
index 896539beaa71..04e479f05fd8 100644
--- a/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py
+++ b/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py
@@ -184,6 +184,8 @@ def test_load_weights_accepts_base_mapper_without_params_map() -> None:
     model = object.__new__(MiniMaxM3ForCausalLM)
     torch.nn.Module.__init__(model)
     model.model_config = model_config
+    # Set by SpecDecOneEngineForCausalLM.__init__, which this test bypasses.
+    model.spec_config = None
     mapper = HfWeightMapper()
     source_name = "model.layers.3.block_sparse_moe.e_score_correction_bias"
     target_name = "model.layers.3.block_sparse_moe.gate.e_score_correction_bias"

From 5d7aadebf11db4ad145b8fe7e7a8774d872eb9ac Mon Sep 17 00:00:00 2001
From: Zheyu Fu <zheyuf@nvidia.com>
Date: Mon, 7 Sep 2026 16:16:45 -0700
Subject: [PATCH 2/7] [TRTLLM-14093][feat] Share the MiniMax-M3 Eagle3 draft KV
 cache with the target manager

MiniMax-M3 keeps its one-model Eagle3 draft layer in the target KV cache
manager (unified draft KV cache) in every supported configuration. The base
manager appends the draft layer after the target's; the M3 manager derives
that layout from the same speculative config and layer mask the base uses.

The draft layer runs the generic TRTLLM attention op, which addresses K/V
through pool pointers, index scales and block offsets and assumes a uniform
per-layer stride within a pool. KVCacheManagerV2 groups sub-pages by size
alone, and at M3's production geometry the per-sparse-layer index-K page is
the same size as a K or V page, so K, V and index-K coalesce into one pool
whose slot is non-uniform (three sub-pages per sparse layer, two per dense or
draft layer). M3's own kernels address their layers through per-layer views
and never notice; the draft layer would be addressed wrongly. 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, with the slot's sub-page count as
the index scale and V one page after K, the same pattern DeepseekV4CacheManager
and the SWA scratch-reuse path use. The draft loop then runs on the shared
manager exactly like every other unified-KV model: no draft-side manager, no
metadata swap, no change to the speculative-decoding code. trtllm-gen accepts
the manager's P128 draft shapes through an opt-in it consults only once its
allowlist has rejected the page size.

Adds the GQA Eagle3 head accuracy test with aggregated and disaggregated
(NIXL) arms plus a chat-GSM8K acceptance probe, and unit tests for the virtual
pool tables, the block-offset copy and the layout derivation.

Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
---
 .../backends/fmha/flashinfer_trtllm_gen.py    |   8 +-
 .../sparse/minimax_m3/cache_manager.py        | 264 ++++++++++++++---
 .../_torch/pyexecutor/py_executor_creator.py  |   4 +-
 .../accuracy/test_disaggregated_serving.py    |   7 +-
 .../defs/accuracy/test_llm_api_pytorch.py     | 279 ++++++++++++++++++
 .../test_lists/qa/llm_function_core.txt       |   2 +
 .../test_lists/test-db/l0_dgx_b200.yml        |   2 +
 .../test_minimax_m3_shared_draft_layers.py    | 127 ++++++++
 8 files changed, 650 insertions(+), 43 deletions(-)
 create mode 100644 tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py

diff --git a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py
index 5c30fea5f7b4..bcc776078ddd 100644
--- a/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py
+++ b/tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py
@@ -770,8 +770,12 @@ def _is_supported_with_reason(
             return False, f"non-positive tokens_per_block ({tokens_per_block})."
         if tokens_per_block & (tokens_per_block - 1) != 0:
             return False, f"tokens_per_block ({tokens_per_block}) that is not a power of 2."
-        if tokens_per_block not in self.SUPPORTED_TOKENS_PER_BLOCK:
-            supported = sorted(self.SUPPORTED_TOKENS_PER_BLOCK)
+        # A KV cache manager may allow extra page sizes, e.g. MiniMax-M3 adds 128.
+        supported_tokens_per_block = self.SUPPORTED_TOKENS_PER_BLOCK | set(
+            getattr(meta.kv_cache_manager, "trtllm_gen_extra_tokens_per_block", ())
+        )
+        if tokens_per_block not in supported_tokens_per_block:
+            supported = sorted(supported_tokens_per_block)
             return False, f"tokens_per_block ({tokens_per_block}). Supported: {supported}."
 
         return True, ""
diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
index 043acaaaa73b..a803b2f377e4 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
@@ -20,12 +20,13 @@
     construction).
   * :class:`MiniMaxM3KVCacheManagerV2` — :class:`KVCacheManagerV2`
     subclass that registers a per-sparse-layer ``Role.INDEX_KEY`` paged
-    buffer alongside the standard K/V buffers.
+    buffer alongside the standard K/V buffers; shared Eagle3 draft layers get
+    their own virtual attention-op pools.
 """
 
 from __future__ import annotations
 
-from typing import List, Optional, Sequence
+from typing import List, Optional, Sequence, Tuple
 
 import torch
 
@@ -39,6 +40,10 @@
 )
 from tensorrt_llm.bindings import DataType
 from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp
+from tensorrt_llm.bindings.internal.batch_manager.kv_cache_manager_v2_utils import (
+    copy_batch_block_offsets_to_device,
+)
+from tensorrt_llm.logger import logger
 from tensorrt_llm.runtime.kv_cache_manager_v2 import BufferConfig, PageIndexMode
 from tensorrt_llm.runtime.kv_cache_manager_v2._common import BAD_PAGE_INDEX
 from tensorrt_llm.runtime.kv_cache_manager_v2._config import DataRole
@@ -133,6 +138,82 @@ def set_index_v(self, layer_idx: int, out_cache_loc: torch.Tensor, idx_v: torch.
         buf.index_copy_(0, out_cache_loc.to(torch.long), idx_v.to(buf.dtype))
 
 
+def shared_draft_layer_count(spec_config, layer_mask) -> int:
+    """How many one-model draft layers the base manager appends to this one.
+
+    Same rule as ``get_pp_layers``: only with a speculative config and no
+    ``layer_mask`` (a separate draft manager setup always passes a mask).
+    """
+    if spec_config is None or layer_mask is not None:
+        return 0
+    from tensorrt_llm._torch.speculative.utils import get_num_spec_layers
+
+    return int(get_num_spec_layers(spec_config))
+
+
+def derive_shared_draft_layout(
+    num_layers: Optional[int],
+    num_kv_heads,
+    num_draft: int,
+) -> Tuple[List[int], Optional[int]]:
+    """Locate the draft layers the base manager appends after the target's.
+
+    ``num_layers`` is the target count; a per-layer ``num_kv_heads`` list (a
+    drafter with a different head count) already includes the draft tail and
+    pins the total. Returns ``(draft_layer_ids, num_target_layers)``, or
+    ``([], None)`` when nothing pins the range.
+    """
+    num_draft = max(0, int(num_draft))
+    if isinstance(num_kv_heads, (list, tuple)):
+        total = len(num_kv_heads)
+        if num_layers is not None:
+            total = max(total, int(num_layers))
+    elif num_layers is not None:
+        total = int(num_layers) + num_draft
+    else:
+        return [], None
+    num_target = total - num_draft
+    return list(range(num_target, total)), num_target
+
+
+def extend_attention_op_pools_for_shared_draft_layers(
+    pool_pointers: torch.Tensor,
+    pool_mapping: torch.Tensor,
+    num_pools: int,
+    draft_layers: Sequence[Tuple[int, int, int]],
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, List[Tuple[int, int]]]:
+    """Give each shared draft layer its own attention-op pool rooted at its K page.
+
+    ``draft_layers`` holds ``(local_layer_idx, key_base_addr, sub_pages_per_slot)``.
+    With ``index_scale = sub_pages_per_slot`` and ``kv_offset = 1``, slot ``s``
+    maps to page ``s * scale`` (K) and ``s * scale + 1`` (V). Returns
+    ``(pool_pointers, pool_mapping, index_scales, kv_offsets, op_pools)`` with
+    ``op_pools = [(attention_op_pool_id, source_storage_pool_id)]``.
+    """
+    pointer_rows = pool_pointers.tolist()
+    mapping_rows = pool_mapping.tolist()
+    nested = pool_pointers.dim() == 3  # NVFP4 carries [data, scale] pointer pairs
+    index_scales: List[int] = []
+    kv_offsets: List[int] = []
+    op_pools: List[Tuple[int, int]] = []
+    for i, (local_layer_idx, key_base_addr, sub_pages_per_slot) in enumerate(draft_layers):
+        op_pool_id = num_pools + i
+        source_pool_id = int(mapping_rows[local_layer_idx][0])
+        pointer_rows.append([[key_base_addr, 0], [0, 0]] if nested else [key_base_addr, 0])
+        mapping_rows[local_layer_idx] = [op_pool_id, 0]
+        index_scales.append(int(sub_pages_per_slot))
+        kv_offsets.append(1)
+        op_pools.append((op_pool_id, source_pool_id))
+    pinned = prefer_pinned()
+    return (
+        torch.tensor(pointer_rows, dtype=pool_pointers.dtype, device="cpu", pin_memory=pinned),
+        torch.tensor(mapping_rows, dtype=pool_mapping.dtype, device="cpu", pin_memory=pinned),
+        torch.tensor(index_scales, dtype=torch.int32, device="cpu", pin_memory=pinned),
+        torch.tensor(kv_offsets, dtype=torch.int32, device="cpu", pin_memory=pinned),
+        op_pools,
+    )
+
+
 class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2):
     """KVCacheManagerV2 subclass with a V2-managed paged index-K cache
     per sparse layer.
@@ -152,9 +233,29 @@ class MiniMaxM3KVCacheManagerV2(KVCacheManagerV2):
       * ``disable_index_value_layer_ids`` — subset whose index-V is
         omitted.
       * ``sparse_index_dim`` — width of the index-K/V vectors.
+
+    Shared Eagle3 draft layers: the base manager appends them after the target
+    layers. They run the generic TRTLLM attention op, which needs a uniform
+    per-layer stride inside a pool. M3's pool is not uniform: here an index-K
+    page is as large as a K or V page, so V2 packs K, V and index-K of all
+    layers into one pool (3 sub-pages per sparse layer, 2 per dense or draft
+    layer). M3's own kernels don't care (they use :meth:`get_buffers`), but the
+    draft layer would read the wrong pages. So each draft layer gets its own
+    virtual attention-op pool rooted at its K page, like
+    ``DeepseekV4CacheManager`` does for SWA layers; see
+    :meth:`_prepare_page_table_tensor`. Can go once V2 keeps index-K in its
+    own pool.
     """
 
     _main_kv_mapper_kind = MapperKind.NHD
+    # Virtual attention-op pools for shared draft layers:
+    # (attention_op_pool_id, source_storage_pool_id) plus their copy parameters.
+    _draft_op_pools: Tuple[Tuple[int, int], ...] = ()
+    _draft_index_scales: Optional[torch.Tensor] = None
+    _draft_kv_offsets: Optional[torch.Tensor] = None
+    # Extra page sizes trtllm-gen may use with this manager (see
+    # FlashInferTrtllmGenFmha); set with the virtual pools.
+    trtllm_gen_extra_tokens_per_block: frozenset = frozenset()
 
     def __init__(
         self,
@@ -170,6 +271,7 @@ def __init__(
         # disable_index_value=True, sparse_index_dim=128). Honoring the
         # executor keyword also makes non-default sparse_index_dim values
         # authoritative for the cache layout instead of falling back to 128.
+        # Peeked (not popped) so the base __init__ still receives them.
         sparse_attention_config = kwargs.get("sparse_attention_config")
         num_layers = kwargs.get("num_layers")
         implementation = getattr(sparse_attention_config, "implementation", "triton")
@@ -184,9 +286,16 @@ def __init__(
             raise ValueError(
                 f"MiniMax M3 sparse_index_dim must be greater than 0, got {sparse_index_dim}."
             )
+        # Shared draft layers sit above the target layers and have no index-K
+        # cache, so the sparse-layer default stops at the target range.
+        self._shared_draft_layer_ids, num_target_layers = derive_shared_draft_layout(
+            num_layers,
+            kwargs.get("num_kv_heads"),
+            shared_draft_layer_count(kwargs.get("spec_config"), kwargs.get("layer_mask")),
+        )
         if sparse_layer_ids is None:
-            if num_layers is not None:
-                sparse_layer_ids = list(range(3, int(num_layers)))
+            if num_target_layers is not None:
+                sparse_layer_ids = list(range(3, num_target_layers))
             else:
                 sparse_layer_ids = []
         if disable_index_value_layer_ids is None:
@@ -232,6 +341,78 @@ def __init__(
                     device=device,
                 )
 
+    def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None:
+        """Base pool tables plus one virtual attention-op pool per shared draft layer.
+
+        See the class docstring for why. The virtual pool reuses the source
+        pool's slot ids; :meth:`copy_batch_block_offsets` scales them.
+        """
+        super()._prepare_page_table_tensor(index_mapper_capacity)
+        draft_layers = [
+            layer_idx
+            for layer_idx in self._shared_draft_layer_ids
+            if layer_idx in self.layer_offsets
+        ]
+        if self.is_draft or not draft_layers:
+            return
+        if self.enable_swa_scratch_reuse:
+            raise NotImplementedError(
+                "MiniMax-M3 shared Eagle3 draft layers do not support SWA scratch reuse."
+            )
+        geometry = []
+        for layer_idx in draft_layers:
+            key_base_addr, _dtype, _num_slots, sub_pages_per_slot, _shape = self._kv_slot_geometry(
+                layer_idx
+            )
+            geometry.append((self.layer_offsets[layer_idx], key_base_addr, sub_pages_per_slot))
+        (
+            self.kv_cache_pool_pointers,
+            self.kv_cache_pool_mapping,
+            self._draft_index_scales,
+            self._draft_kv_offsets,
+            op_pools,
+        ) = extend_attention_op_pools_for_shared_draft_layers(
+            self.kv_cache_pool_pointers, self.kv_cache_pool_mapping, self.num_pools, geometry
+        )
+        self._draft_op_pools = tuple(op_pools)
+        self.num_attention_op_pools = self.num_pools + len(op_pools)
+        # Draft layers run at the target's 128-token pages. trtllm-gen has P128
+        # kernels for their dense-GQA shapes but not for every shape, so opt in
+        # here rather than in the global allowlist.
+        if self.tokens_per_block == 128:
+            self.trtllm_gen_extra_tokens_per_block = frozenset({128})
+        logger.info(
+            f"[unified-kv] draft layers {draft_layers} share the target KV cache manager; "
+            f"attention-op pools {[pool for pool, _ in op_pools]} address their pages."
+        )
+
+    def copy_batch_block_offsets(
+        self,
+        dst_tensor: torch.Tensor,
+        request_ids: List[int],
+        beam_width: int,
+        num_contexts: int,
+        num_seqs: int,
+        max_blocks: Optional[int] = None,
+    ) -> None:
+        super().copy_batch_block_offsets(
+            dst_tensor, request_ids, beam_width, num_contexts, num_seqs, max_blocks=max_blocks
+        )
+        if not self._draft_op_pools:
+            return
+        # Fill each virtual pool from its source pool's slot ids with the draft
+        # layer's scale.
+        copy_idx = self.index_mapper.get_copy_index(request_ids, num_contexts, beam_width)
+        for i, (op_pool_id, source_pool_id) in enumerate(self._draft_op_pools):
+            copy_batch_block_offsets_to_device(
+                self.host_kv_cache_block_offsets[source_pool_id : source_pool_id + 1],
+                dst_tensor[op_pool_id : op_pool_id + 1],
+                copy_idx,
+                self._draft_index_scales[i : i + 1],
+                self._draft_kv_offsets[i : i + 1],
+                self._stream.cuda_stream,
+            )
+
     def _extra_buffers_per_layer(self, *, tokens_per_block):
         """Register a per-sparse-layer ``Role.INDEX_KEY`` :class:`BufferConfig`.
 
@@ -308,25 +489,15 @@ def get_index_v_buffer(self, layer_idx: int) -> Optional[torch.Tensor]:
     def has_index_value(self, layer_idx: int) -> bool:
         return layer_idx in self._index_v_buffers
 
-    def get_buffers(
+    def _kv_slot_geometry(
         self, layer_idx: int, kv_layout: Optional[str] = None
-    ) -> Optional[torch.Tensor]:
-        """Return a paged K+V view with strides spanning the coalesced pool.
+    ) -> Tuple[int, torch.dtype, int, int, List[int]]:
+        """Where a layer's K/V live in the coalesced pool.
 
-        The base :meth:`KVCacheManagerV2.get_buffers` produces a
-        ``[num_pages, kv_factor, ...]`` view with contiguous strides
-        that assume the slot holds exactly one layer's K+V. In M3's
-        pool the slot packs K+V for *all* layers of the group
-        (``scale >= 2 * num_layers_in_group``), so the base view's
-        dim-0 stride does not reach the next slot's K for this layer.
-        (When INDEX_KEY's per-block size coincides with K/V's, it is
-        coalesced into the same pool and contributes to ``scale`` too.)
-
-        The override builds a ``[num_slots, scale, ...]`` view rooted
-        at K's base, then slices ``[:, :2]`` to extract K+V. The slice
-        preserves the dim-0 stride (``scale * page_stride``), so
-        ``view[s, 0/1, ...]`` lands on this layer's K/V at slot ``s``.
-        When omitted, ``kv_layout`` follows the selected sparse backend.
+        Returns ``(addr_key, torch_dtype, num_slots, scale, page_shape)``:
+        ``scale`` sub-pages per slot, this layer's K at sub-page 0 and V at
+        sub-page 1 from ``addr_key``. Used by :meth:`get_buffers` and the draft
+        layers' virtual pools. ``kv_layout`` defaults to the backend's layout.
         """
         if kv_layout is None:
             kv_layout = self._main_kv_layout_name()
@@ -334,7 +505,7 @@ def get_buffers(
             raise ValueError(f"Unsupported kv_layout: {kv_layout}")
         if self.kv_cache_type == CacheTypeCpp.SELFKONLY:
             raise NotImplementedError(
-                "MiniMaxM3KVCacheManagerV2.get_buffers does not support SELFKONLY cache type"
+                "MiniMaxM3KVCacheManagerV2 does not support the SELFKONLY cache type"
             )
 
         layer_offset = self.layer_offsets[layer_idx]
@@ -345,13 +516,13 @@ def get_buffers(
         # V2 always lays V immediately after K within the per-layer
         # contribution to a slot. The slice ``[:, :2]`` depends on this.
         assert addr_key + page_stride_value == addr_value, (
-            f"MiniMaxM3 get_buffers requires addr_K + page_stride "
+            f"MiniMaxM3 requires addr_K + page_stride "
             f"== addr_V (V immediately after K in slot); got "
             f"addr_K={addr_key} page_stride_V={page_stride_value} "
             f"addr_V={addr_value} for layer {layer_idx}."
         )
         assert page_stride_key == page_stride_value, (
-            f"MiniMaxM3 get_buffers requires equal K and V page "
+            f"MiniMaxM3 requires equal K and V page "
             f"strides; got K={page_stride_key} V="
             f"{page_stride_value}."
         )
@@ -378,24 +549,38 @@ def get_buffers(
 
         layer_head_dim = self.head_dim_per_layer[layer_offset]
         num_kv_heads = self.num_kv_heads_per_layer[layer_offset]
+        containers = layer_head_dim // element_per_container
 
         if kv_layout == "NHD":
-            full_slot_shape = [
-                num_slots,
-                scale,
-                self.tokens_per_block,
-                num_kv_heads,
-                layer_head_dim // element_per_container,
-            ]
+            page_shape = [self.tokens_per_block, num_kv_heads, containers]
         else:
-            full_slot_shape = [
-                num_slots,
-                scale,
-                num_kv_heads,
-                self.tokens_per_block,
-                layer_head_dim // element_per_container,
-            ]
+            page_shape = [num_kv_heads, self.tokens_per_block, containers]
+        return addr_key, torch_dtype, num_slots, scale, page_shape
+
+    def get_buffers(
+        self, layer_idx: int, kv_layout: Optional[str] = None
+    ) -> Optional[torch.Tensor]:
+        """Return a paged K+V view with strides spanning the coalesced pool.
 
+        The base :meth:`KVCacheManagerV2.get_buffers` produces a
+        ``[num_pages, kv_factor, ...]`` view with contiguous strides
+        that assume the slot holds exactly one layer's K+V. In M3's
+        pool the slot packs K+V for *all* layers of the group
+        (``scale >= 2 * num_layers_in_group``), so the base view's
+        dim-0 stride does not reach the next slot's K for this layer.
+        (When INDEX_KEY's per-block size coincides with K/V's, it is
+        coalesced into the same pool and contributes to ``scale`` too.)
+
+        The override builds a ``[num_slots, scale, ...]`` view rooted
+        at K's base, then slices ``[:, :2]`` to extract K+V. The slice
+        preserves the dim-0 stride (``scale * page_stride``), so
+        ``view[s, 0/1, ...]`` lands on this layer's K/V at slot ``s``.
+        When omitted, ``kv_layout`` follows the selected sparse backend.
+        """
+        addr_key, torch_dtype, num_slots, scale, page_shape = self._kv_slot_geometry(
+            layer_idx, kv_layout
+        )
+        full_slot_shape = [num_slots, scale, *page_shape]
         full_view = convert_to_torch_tensor(TensorWrapper(addr_key, torch_dtype, full_slot_shape))
         return full_view[:, :2]
 
@@ -504,5 +689,8 @@ def get_minimax_m3_kv_cache_manager_cls():
 __all__ = [
     "MiniMaxM3KVCacheManagerV2",
     "MiniMaxM3SparseIndexCache",
+    "derive_shared_draft_layout",
+    "extend_attention_op_pools_for_shared_draft_layers",
     "get_minimax_m3_kv_cache_manager_cls",
+    "shared_draft_layer_count",
 ]
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index 8cdc9e5e46de..daa520aa781c 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -500,7 +500,9 @@ def create_py_executor(
         # drafters, which it stranded on their private max_seq_len-dense arena.
         is_standalone_drafter = (spec_config.spec_dec_mode.is_dflash()
                                  or spec_config.spec_dec_mode.is_dspark())
-        if cache_transceiver_config is not None and not is_standalone_drafter:
+        # MiniMax-M3 always shares the target KV cache with its Eagle3 drafter.
+        if ((cache_transceiver_config is not None and not is_standalone_drafter)
+                or is_minimax_m3(m3_sparse_config)):
             spec_config._allow_separate_draft_kv_cache = False
 
     # chunk_unit_size may be changed to 64 when using flash mla
diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py
index 03bce01d822b..0c9ae242c222 100644
--- a/tests/integration/defs/accuracy/test_disaggregated_serving.py
+++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py
@@ -62,7 +62,9 @@ def result(self):
         return self
 
 
-DuckLLM = namedtuple('DuckLLM', ['args', 'tokenizer', 'generate_async'])
+DuckLLM = namedtuple('DuckLLM',
+                     ['args', 'tokenizer', 'generate_async', 'router_url'],
+                     defaults=(None, ))
 
 # Timeout for the entire test
 DEFAULT_TEST_TIMEOUT = 3600
@@ -540,7 +542,8 @@ def _show_kvcache_time(kv_cache_perf_dir, max_lines=100):
 
         tokenizer = load_hf_tokenizer(model_name)
         try:
-            yield DuckLLM(args, tokenizer, generate_async)
+            yield DuckLLM(args, tokenizer, generate_async,
+                          f"http://localhost:{serve_port}")
         finally:
             if enable_perf:
                 _show_kvcache_time(kv_cache_perf_dir)
diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
index 41c43a893b7e..ccb31eb94bd4 100644
--- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py
+++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
@@ -16,6 +16,7 @@
 import json
 import os
 import sys
+import time
 from typing import Optional
 from unittest import mock
 
@@ -8913,3 +8914,281 @@ def test_nvfp4(self, use_msa):
             task.evaluate(llm)
             task = GSM8K(model_name)
             task.evaluate(llm)
+
+    def _run_nvfp4_eagle3_disagg(self, model_name, model_path, max_draft_len,
+                                 attention_dp, overlap_scheduler, use_msa,
+                                 cuda_graph):
+        """Disaggregated arm of test_nvfp4_eagle3.
+
+        Context TP2 -> generation TP2 over NIXL on one 4-GPU node, block reuse
+        on the context server. GSM8K through the router plus a chat-GSM8K
+        acceptance probe, since accuracy alone does not notice a corrupted
+        drafter KV (rejected drafts are re-verified).
+        """
+        if not (overlap_scheduler and cuda_graph and use_msa):
+            pytest.skip("the disagg arm pins the production serving shape "
+                        "(overlap scheduler + CUDA graphs + MSA)")
+        import requests
+
+        from .test_disaggregated_serving import launch_disaggregated_llm
+        speculative_config = {
+            "decoding_type": "Eagle3",
+            "max_draft_len": max_draft_len,
+            "speculative_model": f"{llm_models_root()}/MiniMax-M3-EAGLE3-GQA",
+            "eagle3_one_model": True,
+        }
+        common_config = {
+            "speculative_config": speculative_config,
+            "sparse_attention_config": {
+                "algorithm": "minimax_m3",
+                "implementation": "msa",
+                "indexer_kv_dtype": "fp8",
+            },
+            "cache_transceiver_config": {
+                "backend": "NIXL",
+                "transceiver_runtime": "PYTHON",
+                "kv_cache_bounce_size_mb": 0,
+                "kv_transfer_timeout_ms": 600000,
+            },
+            "moe_config": {
+                "backend": "CUTLASS"
+            },
+            "scheduler_config": {
+                "capacity_scheduler_policy": "MAX_UTILIZATION"
+            },
+            "max_seq_len": 4096,
+            "trust_remote_code": True,
+        }
+        kv_cache_common = {
+            "dtype": "fp8",
+            "tokens_per_block": 128,
+            "use_kv_cache_manager_v2": True,
+            "event_buffer_max_size": 0,
+            "free_gpu_memory_fraction": 0.7,
+        }
+        ctx_server_config = {
+            **common_config,
+            "tensor_parallel_size": 2,
+            "moe_expert_parallel_size": 2,
+            "disable_overlap_scheduler": not overlap_scheduler,
+            "enable_attention_dp": False,
+            "enable_chunked_prefill": True,
+            "kv_cache_config": {
+                **kv_cache_common, "enable_block_reuse": True
+            },
+            "max_batch_size": 4,
+            "max_num_tokens": 8192,
+            "cuda_graph_config": None,
+        }
+        gen_server_config = {
+            **common_config,
+            "tensor_parallel_size": 2,
+            "moe_expert_parallel_size": 2,
+            "disable_overlap_scheduler": not overlap_scheduler,
+            "enable_attention_dp": attention_dp,
+            "kv_cache_config": {
+                **kv_cache_common, "enable_block_reuse": False
+            },
+            "max_batch_size": 16,
+            # Decode-only token budget: (1 + draft_len) verify tokens per
+            # request x max_batch_size.
+            "max_num_tokens": (1 + max_draft_len) * 16,
+            "enable_iter_perf_stats": True,
+            # Keep the whole acceptance-probe window: the default buffers
+            # retain only their latest 1000 iterations.
+            "max_stats_len": 5000,
+            "iter_stats_max_iterations": 5000,
+            "cuda_graph_config": {
+                "enable_padding": True,
+                "batch_sizes": [1, 2, 4, 8, 16],
+            },
+        }
+        disaggregated_server_config = {
+            "hostname": "localhost",
+            "backend": "pytorch",
+            "context_servers": {
+                "num_instances": 1
+            },
+            "generation_servers": {
+                "num_instances": 1
+            },
+        }
+        with launch_disaggregated_llm(disaggregated_server_config,
+                                      ctx_server_config,
+                                      gen_server_config,
+                                      model_path,
+                                      server_waiting_timeout=1800) as llm:
+            # The launcher's args only carry the model name (it guesses NVFP4
+            # from it); fill in what this arm actually runs so the same
+            # reference row as the aggregated arm is used.
+            llm.args.quant_config.quant_algo = QuantAlgo.MIXED_PRECISION
+            llm.args.quant_config.kv_cache_quant_algo = QuantAlgo.FP8
+            llm.args.speculative_config = Eagle3DecodingConfig(
+                max_draft_len=max_draft_len,
+                speculative_model=speculative_config["speculative_model"])
+            task = GSM8K(model_name)
+            task.evaluate(llm)
+
+            # Acceptance probe: 200 chat-format GSM8K questions, greedy, 512
+            # tokens, read from the generation worker's /metrics (the disagg
+            # fixture has no get_stats).
+            info = requests.get(f"{llm.router_url}/cluster_info",
+                                timeout=30).json()
+            gen_worker = info["current_workers"]["generation_servers"][0]
+            metrics_url = f'http://{gen_worker["host"]}:{gen_worker["port"]}/metrics'
+            questions = [
+                r["question"]
+                for r in load_dataset(GSM8K.DATASET_DIR, "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
+            ]
+            # Let the background stats collector drain before reading /metrics.
+            time.sleep(1)
+            requests.get(metrics_url, timeout=120).raise_for_status()
+            probe_params = SamplingParams(max_tokens=512, temperature=0)
+            for future in [
+                    llm.generate_async(prompt, probe_params)
+                    for prompt in chat_prompts
+            ]:
+                future.result()
+            time.sleep(1)
+            response = requests.get(metrics_url, timeout=120)
+            response.raise_for_status()
+            drafted = accepted = steps = 0
+            for record in response.json():
+                stats = record.get("specDecodingStats") or {}
+                drafted += stats.get("numDraftTokens", 0)
+                accepted += stats.get("numAcceptedTokens", 0)
+                steps += stats.get("numRequestsWithDraftTokens", 0)
+            assert steps > 0, "no speculative iterations in /metrics"
+            chat_rate = accepted / drafted
+            chat_length = 1 + accepted / steps
+            print(f"MiniMax-M3 Eagle3 disagg chat-GSM8K acceptance: rate="
+                  f"{chat_rate:.3f}, mean acceptance length="
+                  f"{chat_length:.3f} ({steps} spec iterations)")
+            assert chat_rate > 0.80, \
+                f"Eagle3 chat-GSM8K acceptance rate too low: " \
+                f"{chat_rate:.3f} (threshold 0.80, reference 0.839 from " \
+                f"the drafter card)"
+            assert chat_length > 3.4, \
+                f"Eagle3 chat-GSM8K acceptance length too low: " \
+                f"{chat_length:.3f} (threshold 3.4, reference 3.518 from " \
+                f"the drafter card)"
+
+    @pytest.mark.skip_less_device(4)
+    @pytest.mark.skip_less_device_memory(140000)
+    @parametrize_with_ids("disagg", [False, True])
+    @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, disagg):
+        # One-model Eagle3 on the MSA backend; the GQA drafter shares the
+        # target KV cache. MMLU + GSM8K plus a chat-GSM8K acceptance probe,
+        # since accuracy alone does not notice a corrupted drafter KV.
+        if use_msa:
+            from tensorrt_llm._torch.attention.backends.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
+        if disagg:
+            self._run_nvfp4_eagle3_disagg(model_name, model_path, max_draft_len,
+                                          attention_dp, overlap_scheduler,
+                                          use_msa, cuda_graph)
+            return
+        spec_config = Eagle3DecodingConfig(
+            max_draft_len=max_draft_len,
+            speculative_model=f"{llm_models_root()}/MiniMax-M3-EAGLE3-GQA",
+        )
+        # The MSA path runs an FP8 KV cache, as in test_nvfp4.
+        kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.6,
+                                        enable_block_reuse=False,
+                                        dtype="fp8" if use_msa else "auto")
+        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",
+                    indexer_kv_dtype="fp8" if use_msa else "bf16"),
+                moe_config=MoeConfig(backend="CUTLASS"),
+                max_seq_len=4096,
+                # fmha_sm100 caps total_q x heads at 65536; with 4 verify
+                # tokens per row that is 512 (256 with unsharded heads).
+                max_batch_size=256 if attention_dp else 512,
+                speculative_config=spec_config,
+                # CUDA graphs + spec decoding needs the MSA path.
+                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,
+                # Keep the whole acceptance probe: the default buffer holds
+                # only the latest 1000 iterations.
+                max_stats_len=5000,
+                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)
+
+            # Acceptance probe: 200 chat-format GSM8K questions, greedy, 512 tokens.
+            questions = [
+                r["question"]
+                for r in load_dataset(GSM8K.DATASET_DIR, "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
+            # Reference: the MHA drafter card (Inferact/MiniMax-M3-EAGLE3)
+            # reports 0.839 / 3.518; the GQA head measures 0.838 / 3.515 here.
+            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.80, \
+                f"Eagle3 chat-GSM8K acceptance rate too low: {chat_rate:.3f} " \
+                f"(threshold 0.80, reference 0.839 from the drafter card)"
+            assert chat_length > 3.4, \
+                f"Eagle3 chat-GSM8K acceptance length too low: " \
+                f"{chat_length:.3f} (threshold 3.4, reference 3.518 from " \
+                f"the drafter card)"
diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt
index 88927f142ccd..5a659108e65d 100644
--- a/tests/integration/test_lists/qa/llm_function_core.txt
+++ b/tests/integration/test_lists/qa/llm_function_core.txt
@@ -613,6 +613,8 @@ accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph
 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]
 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-disagg=False] 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-disagg=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/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml
index 2f7c7fb8ea9c..05e0a5ddb9d3 100644
--- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml
+++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml
@@ -53,6 +53,8 @@ l0_dgx_b200:
   - disaggregated/test_disaggregated.py::test_disaggregated_gpt_oss_120b_harmony[gpt_oss/gpt-oss-120b]
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[latency_adp_lmtp_tp4]
   - 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-disagg=False] 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-disagg=True] TIMEOUT (180)
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] TIMEOUT (60)
   - unittest/_torch/modeling/test_modeling_deepseekv4.py
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_auto_dtype TIMEOUT (60)
diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py
new file mode 100644
index 000000000000..43bbb1d436a3
--- /dev/null
+++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py
@@ -0,0 +1,127 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Pure-logic tests for MiniMax-M3's shared Eagle3 draft layers: each gets its own
+virtual attention-op pool rooted at its K page inside the mega-slot.
+"""
+
+from types import SimpleNamespace
+
+import torch
+
+from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import (
+    cache_manager as m3_cache_manager,
+)
+from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.cache_manager import (
+    MiniMaxM3KVCacheManagerV2,
+    derive_shared_draft_layout,
+    extend_attention_op_pools_for_shared_draft_layers,
+    shared_draft_layer_count,
+)
+from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2
+
+DRAFT_LOCAL_LAYER = 60
+SCALE = 179  # sub-pages per M3 mega-slot: 3 dense x 2 + 57 sparse x 3 + draft x 2
+DRAFT_K_ADDR = 0x7000_0000
+
+
+def test_virtual_pool_is_rooted_at_the_draft_k_page():
+    pool_pointers = torch.tensor([[0x6000_0000, 0]], dtype=torch.int64)
+    pool_mapping = torch.tensor([[0, i] for i in range(61)], dtype=torch.int32)
+
+    pointers, mapping, index_scales, kv_offsets, op_pools = (
+        extend_attention_op_pools_for_shared_draft_layers(
+            pool_pointers, pool_mapping, 1, [(DRAFT_LOCAL_LAYER, DRAFT_K_ADDR, SCALE)]
+        )
+    )
+
+    assert pointers.tolist() == [[0x6000_0000, 0], [DRAFT_K_ADDR, 0]]
+    # The draft layer moves to the new pool at offset 0; target rows are untouched.
+    assert mapping[DRAFT_LOCAL_LAYER].tolist() == [1, 0]
+    assert mapping[:DRAFT_LOCAL_LAYER].tolist() == pool_mapping[:DRAFT_LOCAL_LAYER].tolist()
+    # Slot s -> page s * SCALE for K and s * SCALE + 1 for V.
+    assert index_scales.tolist() == [SCALE]
+    assert kv_offsets.tolist() == [1]
+    assert op_pools == [(1, 0)]
+
+
+def test_block_offset_copy_fills_the_virtual_pool_from_the_source_pool(monkeypatch):
+    """The base copy fills the storage pools; the override then fills the virtual pool
+    from the source pool's slot ids, and does nothing extra without draft layers.
+    """
+    calls = []
+
+    def fake_base_copy(
+        self, dst_tensor, request_ids, beam_width, num_contexts, num_seqs, max_blocks=None
+    ):
+        calls.append(("base", request_ids, num_seqs, max_blocks))
+
+    def fake_device_copy(host, dst, copy_idx, index_scales, kv_offsets, stream):
+        calls.append(
+            ("virtual", host, dst, copy_idx, index_scales.tolist(), kv_offsets.tolist(), stream)
+        )
+
+    monkeypatch.setattr(KVCacheManagerV2, "copy_batch_block_offsets", fake_base_copy)
+    monkeypatch.setattr(m3_cache_manager, "copy_batch_block_offsets_to_device", fake_device_copy)
+
+    manager = MiniMaxM3KVCacheManagerV2.__new__(MiniMaxM3KVCacheManagerV2)
+    manager._draft_op_pools = ((1, 0),)
+    manager._draft_index_scales = torch.tensor([SCALE], dtype=torch.int32)
+    manager._draft_kv_offsets = torch.tensor([1], dtype=torch.int32)
+    manager.host_kv_cache_block_offsets = torch.zeros((1, 4, 2, 8), dtype=torch.int32)
+    copy_idx = torch.tensor([2, 0], dtype=torch.int32)
+    manager.index_mapper = SimpleNamespace(get_copy_index=lambda ids, nc, bw: copy_idx)
+    manager._stream = SimpleNamespace(cuda_stream=1234)
+    dst = torch.zeros((2, 4, 2, 8), dtype=torch.int32)
+
+    manager.copy_batch_block_offsets(dst, [7, 9], 1, 0, 2, max_blocks=5)
+
+    assert calls[0] == ("base", [7, 9], 2, 5)
+    kind, host, dst_slice, idx, scales, offsets, stream = calls[1]
+    assert kind == "virtual"
+    assert (
+        host.shape == (1, 4, 2, 8)
+        and host.data_ptr() == manager.host_kv_cache_block_offsets.data_ptr()
+    )
+    assert dst_slice.shape == (1, 4, 2, 8) and dst_slice.data_ptr() == dst[1].data_ptr()
+    assert idx is copy_idx
+    assert scales == [SCALE] and offsets == [1] and stream == 1234
+
+    # Non-speculative MiniMax-M3: the base copy is all that runs.
+    calls.clear()
+    plain = MiniMaxM3KVCacheManagerV2.__new__(MiniMaxM3KVCacheManagerV2)
+    plain.copy_batch_block_offsets(dst, [7], 1, 0, 1)
+    assert calls == [("base", [7], 1, None)]
+
+
+def test_draft_layout_locates_the_appended_tail():
+    # num_layers is the target count; a per-layer head list already includes the
+    # draft tail.
+    assert derive_shared_draft_layout(60, 4, 1) == ([60], 60)
+    assert derive_shared_draft_layout(60, [4] * 60 + [64], 1) == ([60], 60)
+    assert derive_shared_draft_layout(60, 4, 0) == ([], 60)
+    assert derive_shared_draft_layout(None, 4, 1) == ([], None)
+
+
+def test_shared_draft_layer_count_follows_the_base_manager():
+    # Same rule as get_pp_layers: a spec config and no layer_mask.
+    mode = SimpleNamespace(
+        is_mtp_eagle_one_model=lambda: False,
+        is_mtp_vanilla=lambda: False,
+        is_eagle3_one_model=lambda: True,
+    )
+    eagle3 = SimpleNamespace(spec_dec_mode=mode, _num_draft_hidden_layers=None)
+    assert shared_draft_layer_count(eagle3, None) == 1
+    assert shared_draft_layer_count(eagle3, [True] * 60) == 0
+    assert shared_draft_layer_count(None, None) == 0

From 9f7305111d72abf265219dc161097e0af2dab959 Mon Sep 17 00:00:00 2001
From: Zheyu Fu <zheyuf@nvidia.com>
Date: Wed, 9 Sep 2026 14:30:20 -0700
Subject: [PATCH 3/7] [TRTLLM-14093][chore] Apply review feedback to the
 MiniMax-M3 Eagle3 port

- Mirror the fmha_sm100 plan's seqused_k in the CUDA-graph-stable plan
  buffers; the planner allocates it per step and the kernel reads it at
  launch, so the captured address must not move.
- Force the shared draft KV cache only for MiniMax-M3 one-model Eagle3, and
  reject the triton reference backend with CUDA graphs: its multi-token
  verify goes through the prefill builder, which cannot be captured. Drop the
  unreachable multi-token dense decode branch and the decode_qo_len field.
- Reject NVFP4 pool pointers in the virtual attention-op pool builder
  instead of writing a null block-scale pointer.
- Size the MSA proxy scratch from the KV cache manager's draft length, so
  runs without speculative decoding size by max_num_sequences.
- Document why per-token seqused_k mirrors the planner (0 for an empty row)
  while the valid-block count is clamped to 1.
- Drop the pinned cuda_graph/use_msa axes of test_nvfp4_eagle3 and list the
  disagg/overlap combinations explicitly; update the test lists.

Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
---
 .../sparse/minimax_m3/cache_manager.py        |  9 +++-
 .../backends/sparse/minimax_m3/msa_backend.py | 20 ++++++---
 .../sparse/minimax_m3/triton_metadata.py      |  2 -
 .../_torch/models/modeling_minimaxm3.py       | 25 +++++------
 .../_torch/pyexecutor/py_executor_creator.py  | 15 ++++++-
 .../defs/accuracy/test_llm_api_pytorch.py     | 44 ++++++++-----------
 .../test_lists/qa/llm_function_core.txt       |  4 +-
 .../test_lists/test-db/l0_dgx_b200.yml        |  4 +-
 .../attention/sparse/msa/test_msa_backend.py  |  5 ++-
 9 files changed, 71 insertions(+), 57 deletions(-)

diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
index a803b2f377e4..d09409e63a50 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
@@ -190,16 +190,21 @@ def extend_attention_op_pools_for_shared_draft_layers(
     ``(pool_pointers, pool_mapping, index_scales, kv_offsets, op_pools)`` with
     ``op_pools = [(attention_op_pool_id, source_storage_pool_id)]``.
     """
+    if pool_pointers.dim() == 3:
+        # NVFP4 pools carry [data, scale] pointer pairs; the draft layer's
+        # block-scale pages would need their own root and stride.
+        raise NotImplementedError(
+            "MiniMax-M3 shared Eagle3 draft layers do not support an NVFP4 KV cache."
+        )
     pointer_rows = pool_pointers.tolist()
     mapping_rows = pool_mapping.tolist()
-    nested = pool_pointers.dim() == 3  # NVFP4 carries [data, scale] pointer pairs
     index_scales: List[int] = []
     kv_offsets: List[int] = []
     op_pools: List[Tuple[int, int]] = []
     for i, (local_layer_idx, key_base_addr, sub_pages_per_slot) in enumerate(draft_layers):
         op_pool_id = num_pools + i
         source_pool_id = int(mapping_rows[local_layer_idx][0])
-        pointer_rows.append([[key_base_addr, 0], [0, 0]] if nested else [key_base_addr, 0])
+        pointer_rows.append([key_base_addr, 0])
         mapping_rows[local_layer_idx] = [op_pool_id, 0]
         index_scales.append(int(sub_pages_per_slot))
         kv_offsets.append(1)
diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py
index c41b4d7c17cb..1fc47723cbfd 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py
@@ -93,7 +93,9 @@ def _worst_case_proxy_max_k_tiles(
 
 
 # Per-step fmha_sm100 plan tensors that must live in CUDA-graph-stable buffers.
-# At num_kv_splits=1 the plan carries no split-KV workspaces, and
+# The planner allocates each of them anew, and the launch reads seqused_k from
+# the plan too, so it needs a stable mirror like the worklists. At
+# num_kv_splits=1 the plan carries no split-KV workspaces, and
 # cute_workspace_buffer is the vendor's cached scratch (kept by reference, not
 # copied).
 _MSA_PLAN_STABLE_KEYS = (
@@ -105,6 +107,7 @@ def _worst_case_proxy_max_k_tiles(
     "qo_segment_lens",
     "kv_segment_lens",
     "qo_offset",
+    "seqused_k",
 )
 _MSA_PLAN_INT64_KEYS = ("packed_work_range", "packed_work_info")
 # fmha_sm100 sizes packed_work_info at 131072 * max(num_kv_splits, 1); forcing
@@ -471,16 +474,20 @@ def _create_msa_buffers(self) -> None:
         self._msa_buffers_ready = True
 
     def _msa_max_decode_tokens(self) -> int:
-        """Worst-case query tokens in one decode step (1 + draft_len per row).
+        """Worst-case query tokens in one decode step: 1 + draft_len per row.
 
+        The KV cache manager knows the speculative config when this metadata is
+        built, so runs without speculative decoding size by max_num_sequences.
         Capped at 16384 so total_q * num_qo_heads stays under fmha_sm100's
         65536 planner limit with 4 index heads.
         """
         max_seqs = int(self.max_num_sequences)
+        draft_len = int(getattr(self.kv_cache_manager, "max_total_draft_tokens", 0) or 0)
+        tokens = max_seqs * (1 + draft_len)
         max_toks = int(self.max_num_tokens or 0)
-        if max_toks <= 0:
-            return max_seqs
-        return max(max_seqs, min(max_toks, 16384))
+        if max_toks > 0:
+            tokens = min(tokens, max_toks)
+        return max(max_seqs, min(tokens, 16384))
 
     def _alloc_msa_proxy_scratch(
         self,
@@ -658,6 +665,9 @@ def on_update_kv_lens(self) -> None:
             "qo_offset": (kv_true - qo_dev).clamp_min(0),
             "seqused_k": kv_true,
         }
+        # seqused_k mirrors the planner (kv length per row, 0 for an empty
+        # row). The clamp_min(1) on n_valid above is a different concern: it
+        # stops top-k from masking every block, which would NaN the GQA row.
         per_token = {
             "kv_segment_lens": kv_true_tok,
             "qo_offset": pos.clamp_min(0),
diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_metadata.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_metadata.py
index 0390e5a158b8..a1bea4f942fc 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_metadata.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/triton_metadata.py
@@ -93,8 +93,6 @@ 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 in decode: 1, or 1 + draft_len with Eagle3.
-    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)
diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py
index 731ef7dbdff3..2495e8dd446b 100644
--- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py
+++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py
@@ -1297,21 +1297,16 @@ def _sdpa_dense_attention_core(
                     )  # [1, H, q, d]
                 output_view[start:end].copy_(out_b.squeeze(0).transpose(0, 1))
         else:
-            # Decode: qo_len query tokens per request; token t of row b attends
-            # seq_lens[b] - qo_len + t + 1 positions (qo_len=1 is the usual mask).
-            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(
+            # Decode: one query token per request at position seq_lens - 1.
+            valid = kv_positions < seq_lens_dev.unsqueeze(-1)  # [batch, max_k]
+            q_b = q_view.view(batch, 1, self.num_heads, self.head_dim).transpose(
                 1, 2
-            )  # [batch, H, qo, d]
-            mask_b = valid.unsqueeze(1)  # [batch, 1, qo, k]
+            )  # [batch, H, 1, d]
+            mask_b = valid.unsqueeze(1).unsqueeze(1)  # [batch, 1, 1, k]
             # Expand K/V one KV head at a time: all heads at once needs an
             # O(batch * max_k * num_heads) temporary that overflows the CUDA-graph
             # pool under attention DP.
-            out_b = q.new_empty(batch, self.num_heads, qo_len, self.head_dim)
+            out_b = q.new_empty(batch, self.num_heads, 1, self.head_dim)
             with sdpa_kernel(_DENSE_SDPA_BACKENDS):
                 for h in range(max(self.num_key_value_heads, 1)):
                     qh = slice(h * group, (h + 1) * group)
@@ -1324,10 +1319,10 @@ def _sdpa_dense_attention_core(
                         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; transpose(1, 2)
-            # .reshape would scramble (head, head_dim) when H != head_dim.
-            output.view(batch, qo_len, self.num_heads, self.head_dim).copy_(out_b.transpose(1, 2))
+                    )  # [batch, group, 1, d]
+            # Drop the singleton query axis; a transpose(1, 2).reshape would
+            # scramble (head, head_dim) when H != head_dim.
+            output.view(batch, self.num_heads, self.head_dim).copy_(out_b.squeeze(2))
 
         return output
 
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index daa520aa781c..95e5647da448 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -500,10 +500,21 @@ def create_py_executor(
         # drafters, which it stranded on their private max_seq_len-dense arena.
         is_standalone_drafter = (spec_config.spec_dec_mode.is_dflash()
                                  or spec_config.spec_dec_mode.is_dspark())
-        # MiniMax-M3 always shares the target KV cache with its Eagle3 drafter.
+        # MiniMax-M3 supports one-model Eagle3 only; its drafter shares the
+        # target KV cache.
+        is_m3_eagle3 = (is_minimax_m3(m3_sparse_config)
+                        and spec_config.spec_dec_mode.is_eagle3_one_model())
         if ((cache_transceiver_config is not None and not is_standalone_drafter)
-                or is_minimax_m3(m3_sparse_config)):
+                or is_m3_eagle3):
             spec_config._allow_separate_draft_kv_cache = False
+        # The triton reference backend runs multi-token verify through its
+        # prefill builder, which cannot be CUDA-graph captured.
+        if (is_m3_eagle3 and m3_sparse_config.implementation != "msa"
+                and llm_args.cuda_graph_config is not None):
+            raise ValueError(
+                "MiniMax-M3 Eagle3 on the triton reference backend does not "
+                "support CUDA graphs; use implementation='msa' or set "
+                "cuda_graph_config=None.")
 
     # chunk_unit_size may be changed to 64 when using flash mla
     attn_runtime_features = AttentionRuntimeFeatures(
diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
index ccb31eb94bd4..a8aff03c95ae 100644
--- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py
+++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
@@ -8916,8 +8916,7 @@ def test_nvfp4(self, use_msa):
             task.evaluate(llm)
 
     def _run_nvfp4_eagle3_disagg(self, model_name, model_path, max_draft_len,
-                                 attention_dp, overlap_scheduler, use_msa,
-                                 cuda_graph):
+                                 attention_dp, overlap_scheduler):
         """Disaggregated arm of test_nvfp4_eagle3.
 
         Context TP2 -> generation TP2 over NIXL on one 4-GPU node, block reuse
@@ -8925,9 +8924,6 @@ def _run_nvfp4_eagle3_disagg(self, model_name, model_path, max_draft_len,
         acceptance probe, since accuracy alone does not notice a corrupted
         drafter KV (rejected drafts are re-verified).
         """
-        if not (overlap_scheduler and cuda_graph and use_msa):
-            pytest.skip("the disagg arm pins the production serving shape "
-                        "(overlap scheduler + CUDA graphs + MSA)")
         import requests
 
         from .test_disaggregated_serving import launch_disaggregated_llm
@@ -9084,29 +9080,29 @@ def _run_nvfp4_eagle3_disagg(self, model_name, model_path, max_draft_len,
 
     @pytest.mark.skip_less_device(4)
     @pytest.mark.skip_less_device_memory(140000)
-    @parametrize_with_ids("disagg", [False, True])
-    @parametrize_with_ids("cuda_graph", [True])
-    @parametrize_with_ids("use_msa", [True])
-    @parametrize_with_ids("overlap_scheduler", [False, True])
+    # The disaggregated arm pins the production serving shape (overlap
+    # scheduler on), so the combinations are listed instead of a full grid.
+    @parametrize_with_ids("disagg,overlap_scheduler", [(False, False),
+                                                       (False, True),
+                                                       (True, 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, disagg):
-        # One-model Eagle3 on the MSA backend; the GQA drafter shares the
-        # target KV cache. MMLU + GSM8K plus a chat-GSM8K acceptance probe,
-        # since accuracy alone does not notice a corrupted drafter KV.
-        if use_msa:
-            from tensorrt_llm._torch.attention.backends.sparse.minimax_m3.msa_utils import \
-                msa_package_available
-            if not msa_package_available():
-                pytest.skip("MSA kernels (fmha_sm100) not available")
+                          overlap_scheduler, disagg):
+        # One-model Eagle3 on the MSA backend with an FP8 KV cache and CUDA
+        # graphs; the GQA drafter shares the target KV cache. MMLU + GSM8K plus
+        # a chat-GSM8K acceptance probe, since accuracy alone does not notice a
+        # corrupted drafter KV.
+        from tensorrt_llm._torch.attention.backends.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
         if disagg:
             self._run_nvfp4_eagle3_disagg(model_name, model_path, max_draft_len,
-                                          attention_dp, overlap_scheduler,
-                                          use_msa, cuda_graph)
+                                          attention_dp, overlap_scheduler)
             return
         spec_config = Eagle3DecodingConfig(
             max_draft_len=max_draft_len,
@@ -9115,26 +9111,24 @@ def test_nvfp4_eagle3(self, tp_size, ep_size, attention_dp,
         # The MSA path runs an FP8 KV cache, as in test_nvfp4.
         kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.6,
                                         enable_block_reuse=False,
-                                        dtype="fp8" if use_msa else "auto")
+                                        dtype="fp8")
         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",
-                    indexer_kv_dtype="fp8" if use_msa else "bf16"),
+                    implementation="msa", indexer_kv_dtype="fp8"),
                 moe_config=MoeConfig(backend="CUTLASS"),
                 max_seq_len=4096,
                 # fmha_sm100 caps total_q x heads at 65536; with 4 verify
                 # tokens per row that is 512 (256 with unsharded heads).
                 max_batch_size=256 if attention_dp else 512,
                 speculative_config=spec_config,
-                # CUDA graphs + spec decoding needs the MSA path.
                 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,
diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt
index 5a659108e65d..ee2f9c545d1f 100644
--- a/tests/integration/test_lists/qa/llm_function_core.txt
+++ b/tests/integration/test_lists/qa/llm_function_core.txt
@@ -613,8 +613,8 @@ accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph
 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]
 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-disagg=False] 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-disagg=True] TIMEOUT (180)
+accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-disagg=False-overlap_scheduler=True] TIMEOUT (180)
+accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=True-disagg=True-overlap_scheduler=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/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml
index 05e0a5ddb9d3..bbbcc63904e6 100644
--- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml
+++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml
@@ -53,8 +53,8 @@ l0_dgx_b200:
   - disaggregated/test_disaggregated.py::test_disaggregated_gpt_oss_120b_harmony[gpt_oss/gpt-oss-120b]
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[latency_adp_lmtp_tp4]
   - 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-disagg=False] 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-disagg=True] TIMEOUT (180)
+  - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-disagg=False-overlap_scheduler=True] TIMEOUT (180)
+  - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=True-disagg=True-overlap_scheduler=True] TIMEOUT (180)
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] TIMEOUT (60)
   - unittest/_torch/modeling/test_modeling_deepseekv4.py
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_auto_dtype TIMEOUT (60)
diff --git a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py
index fbbd9ca0a53f..60a12aeb074a 100644
--- a/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py
+++ b/tests/unittest/_torch/attention/sparse/msa/test_msa_backend.py
@@ -764,12 +764,13 @@ def get_block_ids_per_seq(self, request_ids):
 
 def test_msa_scratch_sizing_covers_spec_verify_tokens():
     """With Eagle3 a decode step has 1 + draft_len query tokens per request, so the
-    proxy scratch must be sized by tokens, not by batch.
+    proxy scratch must be sized by tokens, not by batch. The draft length comes
+    from the KV cache manager, which knows the speculative config at build time.
     """
     metadata_cls = MiniMaxM3MsaSparseAttention.Metadata
     metadata = metadata_cls.__new__(metadata_cls)
-    metadata.kv_cache_manager = None
     # 2 sequences, 4 tokens each (draft_len=3): 8 decode tokens per step.
+    metadata.kv_cache_manager = SimpleNamespace(max_total_draft_tokens=3)
     metadata.max_num_sequences = 2
     metadata.max_num_tokens = 8
     # Store sized for batch-only sizing (4 heads * 16 k-tiles * 2), which is

From 41e15958b4852c355fb239d1c6db7c32800fa9fa Mon Sep 17 00:00:00 2001
From: Zheyu Fu <zheyuf@nvidia.com>
Date: Wed, 9 Sep 2026 15:49:17 -0700
Subject: [PATCH 4/7] [TRTLLM-14093][test] Use a full parameter grid for
 test_nvfp4_eagle3

Parametrize disagg and overlap_scheduler independently instead of listing
combinations; the disaggregated arm no longer skips any of them. Test list
ids follow the new axis order.

Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
---
 tests/integration/defs/accuracy/test_llm_api_pytorch.py | 7 ++-----
 tests/integration/test_lists/qa/llm_function_core.txt   | 4 ++--
 tests/integration/test_lists/test-db/l0_dgx_b200.yml    | 4 ++--
 3 files changed, 6 insertions(+), 9 deletions(-)

diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
index a8aff03c95ae..bc22bbd80dea 100644
--- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py
+++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
@@ -9080,11 +9080,8 @@ def _run_nvfp4_eagle3_disagg(self, model_name, model_path, max_draft_len,
 
     @pytest.mark.skip_less_device(4)
     @pytest.mark.skip_less_device_memory(140000)
-    # The disaggregated arm pins the production serving shape (overlap
-    # scheduler on), so the combinations are listed instead of a full grid.
-    @parametrize_with_ids("disagg,overlap_scheduler", [(False, False),
-                                                       (False, True),
-                                                       (True, True)])
+    @parametrize_with_ids("disagg", [False, 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,
diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt
index ee2f9c545d1f..349b5f37e09e 100644
--- a/tests/integration/test_lists/qa/llm_function_core.txt
+++ b/tests/integration/test_lists/qa/llm_function_core.txt
@@ -613,8 +613,8 @@ accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph
 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]
 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-disagg=False-overlap_scheduler=True] TIMEOUT (180)
-accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=True-disagg=True-overlap_scheduler=True] TIMEOUT (180)
+accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True-disagg=False] TIMEOUT (180)
+accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=True-overlap_scheduler=True-disagg=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/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml
index bbbcc63904e6..028a4985620e 100644
--- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml
+++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml
@@ -53,8 +53,8 @@ l0_dgx_b200:
   - disaggregated/test_disaggregated.py::test_disaggregated_gpt_oss_120b_harmony[gpt_oss/gpt-oss-120b]
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[latency_adp_lmtp_tp4]
   - 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-disagg=False-overlap_scheduler=True] TIMEOUT (180)
-  - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=True-disagg=True-overlap_scheduler=True] TIMEOUT (180)
+  - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True-disagg=False] TIMEOUT (180)
+  - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=True-overlap_scheduler=True-disagg=True] TIMEOUT (180)
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] TIMEOUT (60)
   - unittest/_torch/modeling/test_modeling_deepseekv4.py
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_auto_dtype TIMEOUT (60)

From 8d463129e153500fb2299eaa105bc2cddff5201a Mon Sep 17 00:00:00 2001
From: Zheyu Fu <zheyuf@nvidia.com>
Date: Thu, 10 Sep 2026 15:35:08 -0700
Subject: [PATCH 5/7] [TRTLLM-14093][test] Drop the disagg arm of
 test_nvfp4_eagle3 and clamp decode valid blocks

Review follow-ups on the consolidated MiniMax-M3 Eagle3 PR:

- Remove the disaggregated arm of test_nvfp4_eagle3 (disagg does not take
  new end-to-end accuracy tests); test_disaggregated_serving.py is back to
  main. The grid is now overlap_scheduler x attention_dp.
- CI runs one combination (attention_dp=False, overlap_scheduler=True,
  the production shape) and inherits the stage timeout; the explicit
  TIMEOUT (180) annotations are gone from both lists.
- Clamp the decode-path valid-block counts to at least one block, matching
  the eager path and on_update_kv_lens.

Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
---
 .../backends/sparse/minimax_m3/msa_backend.py |   7 +-
 .../accuracy/test_disaggregated_serving.py    |   7 +-
 .../defs/accuracy/test_llm_api_pytorch.py     | 171 +-----------------
 .../test_lists/qa/llm_function_core.txt       |   3 +-
 .../test_lists/test-db/l0_dgx_b200.yml        |   3 +-
 5 files changed, 10 insertions(+), 181 deletions(-)

diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py
index 1fc47723cbfd..ae1fcb4c2350 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/msa_backend.py
@@ -864,9 +864,12 @@ def _build_step_plans(self) -> None:
         n_valid = per_token_valid_blocks(
             qo_lens_cpu, kv_lens_cpu, qo_offset_cpu, causal=True, block_size=page_size
         )
-        # One entry per query token.
+        # One entry per query token. clamp_min(1) as on the eager path: a
+        # zero-valid row would NaN the GQA row.
         total_q = int(n_valid.shape[0])
-        self.msa_n_valid_blocks[:total_q].copy_(n_valid.to(torch.int32), non_blocking=True)
+        self.msa_n_valid_blocks[:total_q].copy_(
+            n_valid.clamp_min(1).to(torch.int32), non_blocking=True
+        )
 
     def _build_msa_fields(self) -> None:
         """Populate the MSA cache-write buffers for this step.
diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py
index 0c9ae242c222..03bce01d822b 100644
--- a/tests/integration/defs/accuracy/test_disaggregated_serving.py
+++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py
@@ -62,9 +62,7 @@ def result(self):
         return self
 
 
-DuckLLM = namedtuple('DuckLLM',
-                     ['args', 'tokenizer', 'generate_async', 'router_url'],
-                     defaults=(None, ))
+DuckLLM = namedtuple('DuckLLM', ['args', 'tokenizer', 'generate_async'])
 
 # Timeout for the entire test
 DEFAULT_TEST_TIMEOUT = 3600
@@ -542,8 +540,7 @@ def _show_kvcache_time(kv_cache_perf_dir, max_lines=100):
 
         tokenizer = load_hf_tokenizer(model_name)
         try:
-            yield DuckLLM(args, tokenizer, generate_async,
-                          f"http://localhost:{serve_port}")
+            yield DuckLLM(args, tokenizer, generate_async)
         finally:
             if enable_perf:
                 _show_kvcache_time(kv_cache_perf_dir)
diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
index bc22bbd80dea..d6f909957e62 100644
--- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py
+++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py
@@ -16,7 +16,6 @@
 import json
 import os
 import sys
-import time
 from typing import Optional
 from unittest import mock
 
@@ -8915,177 +8914,13 @@ def test_nvfp4(self, use_msa):
             task = GSM8K(model_name)
             task.evaluate(llm)
 
-    def _run_nvfp4_eagle3_disagg(self, model_name, model_path, max_draft_len,
-                                 attention_dp, overlap_scheduler):
-        """Disaggregated arm of test_nvfp4_eagle3.
-
-        Context TP2 -> generation TP2 over NIXL on one 4-GPU node, block reuse
-        on the context server. GSM8K through the router plus a chat-GSM8K
-        acceptance probe, since accuracy alone does not notice a corrupted
-        drafter KV (rejected drafts are re-verified).
-        """
-        import requests
-
-        from .test_disaggregated_serving import launch_disaggregated_llm
-        speculative_config = {
-            "decoding_type": "Eagle3",
-            "max_draft_len": max_draft_len,
-            "speculative_model": f"{llm_models_root()}/MiniMax-M3-EAGLE3-GQA",
-            "eagle3_one_model": True,
-        }
-        common_config = {
-            "speculative_config": speculative_config,
-            "sparse_attention_config": {
-                "algorithm": "minimax_m3",
-                "implementation": "msa",
-                "indexer_kv_dtype": "fp8",
-            },
-            "cache_transceiver_config": {
-                "backend": "NIXL",
-                "transceiver_runtime": "PYTHON",
-                "kv_cache_bounce_size_mb": 0,
-                "kv_transfer_timeout_ms": 600000,
-            },
-            "moe_config": {
-                "backend": "CUTLASS"
-            },
-            "scheduler_config": {
-                "capacity_scheduler_policy": "MAX_UTILIZATION"
-            },
-            "max_seq_len": 4096,
-            "trust_remote_code": True,
-        }
-        kv_cache_common = {
-            "dtype": "fp8",
-            "tokens_per_block": 128,
-            "use_kv_cache_manager_v2": True,
-            "event_buffer_max_size": 0,
-            "free_gpu_memory_fraction": 0.7,
-        }
-        ctx_server_config = {
-            **common_config,
-            "tensor_parallel_size": 2,
-            "moe_expert_parallel_size": 2,
-            "disable_overlap_scheduler": not overlap_scheduler,
-            "enable_attention_dp": False,
-            "enable_chunked_prefill": True,
-            "kv_cache_config": {
-                **kv_cache_common, "enable_block_reuse": True
-            },
-            "max_batch_size": 4,
-            "max_num_tokens": 8192,
-            "cuda_graph_config": None,
-        }
-        gen_server_config = {
-            **common_config,
-            "tensor_parallel_size": 2,
-            "moe_expert_parallel_size": 2,
-            "disable_overlap_scheduler": not overlap_scheduler,
-            "enable_attention_dp": attention_dp,
-            "kv_cache_config": {
-                **kv_cache_common, "enable_block_reuse": False
-            },
-            "max_batch_size": 16,
-            # Decode-only token budget: (1 + draft_len) verify tokens per
-            # request x max_batch_size.
-            "max_num_tokens": (1 + max_draft_len) * 16,
-            "enable_iter_perf_stats": True,
-            # Keep the whole acceptance-probe window: the default buffers
-            # retain only their latest 1000 iterations.
-            "max_stats_len": 5000,
-            "iter_stats_max_iterations": 5000,
-            "cuda_graph_config": {
-                "enable_padding": True,
-                "batch_sizes": [1, 2, 4, 8, 16],
-            },
-        }
-        disaggregated_server_config = {
-            "hostname": "localhost",
-            "backend": "pytorch",
-            "context_servers": {
-                "num_instances": 1
-            },
-            "generation_servers": {
-                "num_instances": 1
-            },
-        }
-        with launch_disaggregated_llm(disaggregated_server_config,
-                                      ctx_server_config,
-                                      gen_server_config,
-                                      model_path,
-                                      server_waiting_timeout=1800) as llm:
-            # The launcher's args only carry the model name (it guesses NVFP4
-            # from it); fill in what this arm actually runs so the same
-            # reference row as the aggregated arm is used.
-            llm.args.quant_config.quant_algo = QuantAlgo.MIXED_PRECISION
-            llm.args.quant_config.kv_cache_quant_algo = QuantAlgo.FP8
-            llm.args.speculative_config = Eagle3DecodingConfig(
-                max_draft_len=max_draft_len,
-                speculative_model=speculative_config["speculative_model"])
-            task = GSM8K(model_name)
-            task.evaluate(llm)
-
-            # Acceptance probe: 200 chat-format GSM8K questions, greedy, 512
-            # tokens, read from the generation worker's /metrics (the disagg
-            # fixture has no get_stats).
-            info = requests.get(f"{llm.router_url}/cluster_info",
-                                timeout=30).json()
-            gen_worker = info["current_workers"]["generation_servers"][0]
-            metrics_url = f'http://{gen_worker["host"]}:{gen_worker["port"]}/metrics'
-            questions = [
-                r["question"]
-                for r in load_dataset(GSM8K.DATASET_DIR, "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
-            ]
-            # Let the background stats collector drain before reading /metrics.
-            time.sleep(1)
-            requests.get(metrics_url, timeout=120).raise_for_status()
-            probe_params = SamplingParams(max_tokens=512, temperature=0)
-            for future in [
-                    llm.generate_async(prompt, probe_params)
-                    for prompt in chat_prompts
-            ]:
-                future.result()
-            time.sleep(1)
-            response = requests.get(metrics_url, timeout=120)
-            response.raise_for_status()
-            drafted = accepted = steps = 0
-            for record in response.json():
-                stats = record.get("specDecodingStats") or {}
-                drafted += stats.get("numDraftTokens", 0)
-                accepted += stats.get("numAcceptedTokens", 0)
-                steps += stats.get("numRequestsWithDraftTokens", 0)
-            assert steps > 0, "no speculative iterations in /metrics"
-            chat_rate = accepted / drafted
-            chat_length = 1 + accepted / steps
-            print(f"MiniMax-M3 Eagle3 disagg chat-GSM8K acceptance: rate="
-                  f"{chat_rate:.3f}, mean acceptance length="
-                  f"{chat_length:.3f} ({steps} spec iterations)")
-            assert chat_rate > 0.80, \
-                f"Eagle3 chat-GSM8K acceptance rate too low: " \
-                f"{chat_rate:.3f} (threshold 0.80, reference 0.839 from " \
-                f"the drafter card)"
-            assert chat_length > 3.4, \
-                f"Eagle3 chat-GSM8K acceptance length too low: " \
-                f"{chat_length:.3f} (threshold 3.4, reference 3.518 from " \
-                f"the drafter card)"
-
     @pytest.mark.skip_less_device(4)
     @pytest.mark.skip_less_device_memory(140000)
-    @parametrize_with_ids("disagg", [False, 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, disagg):
+                          overlap_scheduler):
         # One-model Eagle3 on the MSA backend with an FP8 KV cache and CUDA
         # graphs; the GQA drafter shares the target KV cache. MMLU + GSM8K plus
         # a chat-GSM8K acceptance probe, since accuracy alone does not notice a
@@ -9097,10 +8932,6 @@ def test_nvfp4_eagle3(self, tp_size, ep_size, attention_dp,
         model_name = "nvidia/MiniMax-M3-NVFP4"
         model_path = f"{llm_models_root()}/MiniMax-M3-NVFP4"
         max_draft_len = 3
-        if disagg:
-            self._run_nvfp4_eagle3_disagg(model_name, model_path, max_draft_len,
-                                          attention_dp, overlap_scheduler)
-            return
         spec_config = Eagle3DecodingConfig(
             max_draft_len=max_draft_len,
             speculative_model=f"{llm_models_root()}/MiniMax-M3-EAGLE3-GQA",
diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt
index 349b5f37e09e..4fa8af4ecb19 100644
--- a/tests/integration/test_lists/qa/llm_function_core.txt
+++ b/tests/integration/test_lists/qa/llm_function_core.txt
@@ -613,8 +613,7 @@ accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8_piecewise_cuda_graph
 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]
 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-disagg=False] TIMEOUT (180)
-accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=True-overlap_scheduler=True-disagg=True] TIMEOUT (180)
+accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True]
 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/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml
index 028a4985620e..1f3af15aca99 100644
--- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml
+++ b/tests/integration/test_lists/test-db/l0_dgx_b200.yml
@@ -53,8 +53,7 @@ l0_dgx_b200:
   - disaggregated/test_disaggregated.py::test_disaggregated_gpt_oss_120b_harmony[gpt_oss/gpt-oss-120b]
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus[latency_adp_lmtp_tp4]
   - 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-disagg=False] TIMEOUT (180)
-  - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=True-overlap_scheduler=True-disagg=True] TIMEOUT (180)
+  - accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4_eagle3[tp_size=4-ep_size=4-attention_dp=False-overlap_scheduler=True]
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_4gpus[pp4-mtp_nextn=0-attention_dp=False-cuda_graph=False-overlap_scheduler=False-torch_compile=False] TIMEOUT (60)
   - unittest/_torch/modeling/test_modeling_deepseekv4.py
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV4Flash::test_auto_dtype TIMEOUT (60)

From 0d9bb7da12a617bd41cdd2e5619de0079635d998 Mon Sep 17 00:00:00 2001
From: Zheyu Fu <zheyuf@nvidia.com>
Date: Sun, 13 Sep 2026 23:17:12 -0700
Subject: [PATCH 6/7] [TRTLLM-14093][fix] Skip the virtual draft pools under
 per-layer page tables

A drafter with a different K/V page stride puts KVCacheManagerV2 into
per-layer page tables, where every layer already has its own attention-op
pool and pool ids are layer ids. Appending pools at num_pools + i there
aliased real layers' pools and indexed the host block-offset table with a
layer id.

Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
---
 .../sparse/minimax_m3/cache_manager.py        | 13 +++++----
 .../test_minimax_m3_shared_draft_layers.py    | 28 +++++++++++++++++++
 2 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
index db2cd5233bab..0d6adae238ac 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
@@ -364,6 +364,14 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None:
             raise NotImplementedError(
                 "MiniMax-M3 shared Eagle3 draft layers do not support SWA scratch reuse."
             )
+        # Draft layers run at the target's 128-token pages. trtllm-gen has P128
+        # kernels for their dense-GQA shapes but not for every shape, so opt in
+        # here rather than in the global allowlist.
+        if self.tokens_per_block == 128:
+            self.trtllm_gen_extra_tokens_per_block = frozenset({128})
+        if self._use_per_layer_page_tables:
+            # Per-layer page tables already give every layer its own pool.
+            return
         geometry = []
         for layer_idx in draft_layers:
             key_base_addr, _dtype, _num_slots, sub_pages_per_slot, _shape = self._kv_slot_geometry(
@@ -381,11 +389,6 @@ def _prepare_page_table_tensor(self, index_mapper_capacity: int) -> None:
         )
         self._draft_op_pools = tuple(op_pools)
         self.num_attention_op_pools = self.num_pools + len(op_pools)
-        # Draft layers run at the target's 128-token pages. trtllm-gen has P128
-        # kernels for their dense-GQA shapes but not for every shape, so opt in
-        # here rather than in the global allowlist.
-        if self.tokens_per_block == 128:
-            self.trtllm_gen_extra_tokens_per_block = frozenset({128})
         logger.info(
             f"[unified-kv] draft layers {draft_layers} share the target KV cache manager; "
             f"attention-op pools {[pool for pool, _ in op_pools]} address their pages."
diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py
index 43bbb1d436a3..4e7b2e29487e 100644
--- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py
+++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py
@@ -105,6 +105,34 @@ def fake_device_copy(host, dst, copy_idx, index_scales, kv_offsets, stream):
     assert calls == [("base", [7], 1, None)]
 
 
+def test_per_layer_page_tables_get_no_virtual_pools(monkeypatch):
+    """With per-layer page tables every layer already has its own pool."""
+    pointers = torch.tensor([[0x6000_0000 + i, 0] for i in range(61)])
+    mapping = torch.tensor([[i, 0] for i in range(61)], dtype=torch.int32)
+
+    def fake_base_prepare(self, index_mapper_capacity):
+        self._use_per_layer_page_tables = True
+        self.kv_cache_pool_pointers = pointers.clone()
+        self.kv_cache_pool_mapping = mapping.clone()
+        self.num_attention_op_pools = 61
+
+    monkeypatch.setattr(KVCacheManagerV2, "_prepare_page_table_tensor", fake_base_prepare)
+    manager = MiniMaxM3KVCacheManagerV2.__new__(MiniMaxM3KVCacheManagerV2)
+    manager._shared_draft_layer_ids = [DRAFT_LOCAL_LAYER]
+    manager.layer_offsets = {i: i for i in range(61)}
+    manager.is_draft = False
+    manager.enable_swa_scratch_reuse = False
+    manager.tokens_per_block = 128
+
+    manager._prepare_page_table_tensor(8)
+
+    assert torch.equal(manager.kv_cache_pool_pointers, pointers)
+    assert torch.equal(manager.kv_cache_pool_mapping, mapping)
+    assert manager.num_attention_op_pools == 61
+    assert manager._draft_op_pools == ()
+    assert manager.trtllm_gen_extra_tokens_per_block == frozenset({128})
+
+
 def test_draft_layout_locates_the_appended_tail():
     # num_layers is the target count; a per-layer head list already includes the
     # draft tail.

From 73576eddb1bd35af7cd069d47245562a3eaa17ea Mon Sep 17 00:00:00 2001
From: Zheyu Fu <zheyuf@nvidia.com>
Date: Mon, 14 Sep 2026 00:56:28 -0700
Subject: [PATCH 7/7] [TRTLLM-14093][fix] Refuse tree decoding for MiniMax-M3
 Eagle3

Only tree acceptance relocates accepted draft KV, and neither the MSA
decode kernels nor M3's coalesced pool layout support it. Reject tree
configs at startup and any relocation request in the manager.

Signed-off-by: Zheyu Fu <zheyuf@nvidia.com>
---
 .../sparse/minimax_m3/cache_manager.py        |  6 ++++++
 .../_torch/pyexecutor/py_executor_creator.py  |  3 +++
 .../test_minimax_m3_shared_draft_layers.py    | 20 +++++++++++++++++++
 3 files changed, 29 insertions(+)

diff --git a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
index 0d6adae238ac..4c52f9e677c6 100644
--- a/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
+++ b/tensorrt_llm/_torch/attention/backends/sparse/minimax_m3/cache_manager.py
@@ -421,6 +421,12 @@ def copy_batch_block_offsets(
                 self._stream.cuda_stream,
             )
 
+    def update_resources(self, scheduled_batch, attn_metadata=None, kv_cache_dtype_byte_size=None):
+        # Only tree acceptance relocates draft KV, which M3's pool layout cannot do.
+        if any(r.py_num_accepted_draft_tokens_indices for r in scheduled_batch.generation_requests):
+            raise NotImplementedError("MiniMax-M3 does not relocate accepted draft tokens.")
+        super().update_resources(scheduled_batch, attn_metadata, kv_cache_dtype_byte_size)
+
     def _extra_buffers_per_layer(self, *, tokens_per_block):
         """Register a per-sparse-layer ``Role.INDEX_KEY`` :class:`BufferConfig`.
 
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
index b08f105c307b..517eb5072d31 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py
@@ -517,6 +517,9 @@ def _create_py_executor_impl(
                 "MiniMax-M3 Eagle3 on the triton reference backend does not "
                 "support CUDA graphs; use implementation='msa' or set "
                 "cuda_graph_config=None.")
+        if is_m3_eagle3 and not spec_config.is_linear_tree:
+            raise ValueError(
+                "MiniMax-M3 Eagle3 supports the linear draft chain only.")
 
     # chunk_unit_size may be changed to 64 when using flash mla
     attn_runtime_features = AttentionRuntimeFeatures(
diff --git a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py
index 4e7b2e29487e..ceda71f86b41 100644
--- a/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py
+++ b/tests/unittest/_torch/attention/sparse/msa/test_minimax_m3_shared_draft_layers.py
@@ -18,6 +18,7 @@
 
 from types import SimpleNamespace
 
+import pytest
 import torch
 
 from tensorrt_llm._torch.attention.backends.sparse.minimax_m3 import (
@@ -133,6 +134,25 @@ def fake_base_prepare(self, index_mapper_capacity):
     assert manager.trtllm_gen_extra_tokens_per_block == frozenset({128})
 
 
+def test_update_resources_refuses_tree_relocation(monkeypatch):
+    """Linear acceptance rewinds through the base; tree acceptance is refused."""
+    calls = []
+    monkeypatch.setattr(KVCacheManagerV2, "update_resources", lambda self, *a: calls.append(a))
+    manager = MiniMaxM3KVCacheManagerV2.__new__(MiniMaxM3KVCacheManagerV2)
+    linear = SimpleNamespace(
+        py_num_accepted_draft_tokens=2, py_num_accepted_draft_tokens_indices=[]
+    )
+    tree = SimpleNamespace(
+        py_num_accepted_draft_tokens=2, py_num_accepted_draft_tokens_indices=[0, 2]
+    )
+
+    manager.update_resources(SimpleNamespace(generation_requests=[linear]), None, 2)
+    assert len(calls) == 1
+
+    with pytest.raises(NotImplementedError, match="relocate"):
+        manager.update_resources(SimpleNamespace(generation_requests=[linear, tree]), None, 2)
+
+
 def test_draft_layout_locates_the_appended_tail():
     # num_layers is the target count; a per-layer head list already includes the
     # draft tail.