[TRTLLM-14093][feat] Eagle3 support for MiniMax-M3

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              | Yes             |
 | `Gemma4UnifiedForConditionalGeneration` | Untested          | Untested   | Untested                   | No                    | Yes             | No                         | Yes           | Untested         | No             | Yes                      | Untested              | Yes             |
 | `Step3p7ForConditionalGeneration`| Yes               | Yes        | Yes                        | Untested              | Untested        | MTP                        | Yes           | Untested         | Untested       | Yes                      | Untested              | Yes             |
-| `MiniMaxM3SparseForConditionalGeneration` [^12] | Yes               | Yes        | Yes                        | Untested              | Untested        | No                         | Yes           | Untested         | No             | N/A                      | Untested              | Yes             |
+| `MiniMaxM3SparseForConditionalGeneration` [^12] | Yes               | Yes        | Yes                        | Untested              | Untested        | EAGLE-3 (Linear)           | Yes           | Untested         | No             | N/A                      | Untested              | Yes             |
 
 [^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.

tensorrt_llm/_torch/attention/backends/fmha/flashinfer_trtllm_gen.py

@@ -760,8 +760,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, ""

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
 
 
@@ -165,7 +165,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 +372,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 +576,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,51 +592,41 @@ 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(
                     f"static_buffers max_num_tokens={static_buffers['max_num_tokens']} "
                     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,24 +641,18 @@ 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(
                     f"static_buffers max_num_tokens={static_buffers['max_num_tokens']} "
                     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 +664,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 +825,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 +865,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",
 ]