[MiniMax-M3.1] Enable Mono decode with fixed DSpark verify

atom/models/minimax_m3/model.py

@@ -856,6 +856,9 @@ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
     def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
         self.model.aux_hidden_state_layers = layers
 
+    def bind_mono_drafter(self, drafter) -> None:
+        self._mono.bind_drafter(drafter)
+
     def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
         """Default Eagle3 aux hidden-state layer ids: early / middle / late of
         the target model (early=2, mid=n//2, late=n-3), matching vLLM's default.
@@ -940,6 +943,9 @@ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
     def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
         self.language_model.set_aux_hidden_state_layers(layers)
 
+    def bind_mono_drafter(self, drafter) -> None:
+        self.language_model.bind_mono_drafter(drafter)
+
     def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
         return self.language_model.get_eagle3_aux_hidden_state_layers()
 

atom/models/minimax_m3/mono/check.py

@@ -3,8 +3,9 @@
 """Layer-by-layer comparison of the mono path against the original modules.
 
 Debug only (``ATOM_MONO_CHECK=1``, meant for ``--enforce-eager``). The dense
-layers run twice, mono then the original modules, compared where they hand the
-first sparse layer its (ar, res); the step goes on from mono's. Every sparse
+layers run twice where mono has kernels for them, mono then the original
+modules, compared where they hand the first sparse layer its (ar, res); the
+step goes on from mono's. Every sparse
 layer runs twice on the same (ar, res): first the mono layer kernel, then the
 original decoder layer, fed ``ar`` on rank 0 and zeros elsewhere so its fused
 all-reduce reproduces ``ar`` exactly. Both write the same cache bytes for the
@@ -36,23 +37,28 @@ def forward_checked(
     runner.mailboxes.begin_step()
     ars = runner.ars
     # the dense layers, mono first, then the original ones on the same step,
-    # compared where they hand the first sparse layer its (ar, res)
-    res = runner.run_dense_mono(input_ids, positions, fwd, (), []).clone()
-    n = res.shape[0]
-    ar = ars[0][:n].clone()
-    ref_res = runner.run_dense(input_ids, positions)
-    if runner.rank == 0:
-        for k in range(n):
-            logger.info(
-                "mono check dense layers token %d/%d pos %d: residual %s | "
-                "reduced out %s",
-                k,
-                n,
-                int(positions[k]),
-                compare(res[k], ref_res[k]),
-                compare(ar[k], ars[0][k]),
-            )
-    ars[0][:n].copy_(ar)
+    # compared where they hand the first sparse layer its (ar, res). A model
+    # with no kernels for its prefix (M3.1) has one arm and nothing to compare.
+    if not runner.dense_mono:
+        res = runner.run_dense(input_ids, positions, (), [])
+        n = res.shape[0]
+    else:
+        res = runner.run_dense_mono(input_ids, positions, fwd, (), []).clone()
+        n = res.shape[0]
+        ar = ars[0][:n].clone()
+        ref_res = runner.run_dense(input_ids, positions, (), [])
+        if runner.rank == 0:
+            for k in range(n):
+                logger.info(
+                    "mono check dense layers token %d/%d pos %d: residual %s | "
+                    "reduced out %s",
+                    k,
+                    n,
+                    int(positions[k]),
+                    compare(res[k], ref_res[k]),
+                    compare(ar[k], ars[0][k]),
+                )
+        ars[0][:n].copy_(ar)
     rows = token_rows(fwd, n)
     zeros = torch.zeros(n, HIDDEN, dtype=ars[0].dtype, device=ars[0].device)
     for i, lw in enumerate(runner.sparse):

atom/models/minimax_m3/mono/dispatch.py

@@ -41,6 +41,13 @@ def _uses_dspark(atom_config) -> bool:
     return bool(use_dspark is not None and use_dspark())
 
 
+def _dspark_mode_supported(atom_config) -> bool:
+    if not _uses_dspark(atom_config):
+        return True
+    dspark = getattr(atom_config, "dspark", None)
+    return dspark is not None and not (dspark.confidence_schedule or dspark.ragged)
+
+
 def _config_refusal(atom_config) -> str | None:
     """Why this deployment cannot use mono, or None."""
     checks = common_refusals(atom_config, Dims, SUPPORTED_TP) + [
@@ -57,9 +64,10 @@ def _config_refusal(atom_config) -> str | None:
             atom_config.max_model_len <= MAX_CONTEXT,
             f"max_model_len {atom_config.max_model_len} > {MAX_CONTEXT}",
         ),
-        # DSpark reads the target's aux hidden states through forward hooks on
-        # the decoder layers, which the mono kernels replace.
-        (not _uses_dspark(atom_config), "DSpark speculative decoding"),
+        (
+            _dspark_mode_supported(atom_config),
+            "DSpark confidence schedule / ragged",
+        ),
     ]
     return first_refusal(checks)
 
@@ -69,9 +77,16 @@ class MonoDecode:
 
     def __init__(self, causal_lm, atom_config) -> None:
         self._lm = causal_lm
+        self._uses_dspark = _uses_dspark(atom_config)
+        self._drafter = None
         # created on its first eligible step (``LazyRunner``)
         self._mono = LazyRunner(
-            lambda: MonoDecodeRunner(causal_lm), "MiniMax-M3 mono decode"
+            lambda: MonoDecodeRunner(
+                causal_lm,
+                drafter=self._drafter,
+                require_aux=self._uses_dspark,
+            ),
+            "MiniMax-M3 mono decode",
         )
         if not envs.ATOM_MONO_ENABLE:
             self._mono.enabled = False
@@ -81,6 +96,13 @@ def __init__(self, causal_lm, atom_config) -> None:
             logger.info("MiniMax-M3 mono decode off: %s", why)
             self._mono.enabled = False
 
+    def bind_drafter(self, drafter) -> None:
+        if not self._uses_dspark:
+            return
+        if self._mono.runner is not None:
+            raise RuntimeError("cannot bind a DSpark drafter after Mono starts")
+        self._drafter = drafter
+
     def supports(
         self,
         input_ids: torch.Tensor,
@@ -151,3 +173,12 @@ def forward(self, input_ids: torch.Tensor, positions: torch.Tensor) -> torch.Ten
                 envs.ATOM_MONO_TRACE, sums,
             )  # fmt: skip
         return out
+
+
+def install_mono_decode(
+    model: torch.nn.Module, _atom_config, drafter, *, metadata_builder=None
+):
+    """Bind MiniMax's in-model Mono router after drafter hooks are armed."""
+    if drafter is not None:
+        model.bind_mono_drafter(drafter)
+    return model

atom/models/minimax_m3/mono/runner.py

@@ -5,8 +5,11 @@
 
 Each dense layer 0..2 is three launches: K1 built dense (``dense_pre``: q, the
 K / V insert), its original attention kernel, then ``dense_post`` (o_proj
-through the MLP, both all-reduces in-kernel -> (ar, h_mid)). Every sparse MoE
-layer is one launch of K4 with the layer's K1 fused in front of its stages::
+through the MLP, both all-reduces in-kernel -> (ar, h_mid)). A model whose
+prefix layers are sparse-attention instead (M3.1) has no kernels for them and
+runs them on the original modules (``dense_prefix_mono``, ``run_dense``), mono
+taking over at the first MoE layer. Every sparse MoE layer is one launch of K4
+with the layer's K1 fused in front of its stages::
 
     K1 part        (ar, res) -> h, q; K / V / index cache insert; indexer scores
     K4 part        top-k, attention .. FFN, both all-reduces in-kernel
@@ -69,6 +72,7 @@
 from atom.models.minimax_m3.mono.layout import SCRATCH_BYTES as K4_SCRATCH_BYTES
 from atom.models.minimax_m3.mono.weights import DenseLayer, SparseMoeLayer
 from atom.mono.plan.execution import BLOCKS
+from atom.mono.runtime.aux_taps import bind_aux_taps
 from atom.mono.runtime.consensus import MonoUnsupported, bind_agreed
 from atom.mono.runtime.debug import (
     DIAG_BYTES,
@@ -121,21 +125,36 @@ def token_rows(fwd, n: int) -> tuple[torch.Tensor, torch.Tensor]:
     )
 
 
+def dense_prefix_mono(layers) -> bool:
+    """Whether layers 0..N_DENSE-1 run on the fused dense kernels.
+
+    The prefix never runs on the MoE kernels, so its FFN has to be the dense
+    MLP. M3's prefix is dense full-attention and takes ``dense_pre`` /
+    ``dense_post``; M3.1's is sparse attention + indexer, which only the
+    original modules serve (``run_dense``)."""
+    if len(layers) <= N_DENSE or any(layers[i].is_moe_layer for i in range(N_DENSE)):
+        raise MonoUnsupported(f"layers 0..{N_DENSE - 1} must be dense-MLP layers")
+    return not any(
+        getattr(layers[i].self_attn, "is_indexed_sparse_attention", False)
+        for i in range(N_DENSE)
+    )
+
+
 class MonoDecodeRunner:
     """Owns the per-layer weight views, scratch, peer buffers and compiled kernels.
 
     Construction is collective over the TP group: every rank binds and checks its
     shard, then all agree (``consensus.bind_agreed``) before the peer handshake, so a rank that
     refuses turns mono off on every rank instead of leaving the others waiting."""
 
-    def __init__(self, causal_lm) -> None:
+    def __init__(self, causal_lm, drafter=None, require_aux: bool = False) -> None:
         group = get_tp_group().cpu_group
-        bind_agreed(lambda: self._bind(causal_lm), group)
+        bind_agreed(lambda: self._bind(causal_lm, drafter, require_aux), group)
         self._allocate(group)
         # per decode batch size (every graph captures its own) a sparse layer
-        # kernel of each kind the layers use (``self.k4[n][lw.index_topk]``) and
-        # the dense layers' (``self.k4[n]["dense_pre0" | "dense_pre" |
-        # "dense_post"]``), built and compiled on its first step (``prepare``)
+        # kernel of each kind the layers use (``self.k4[n][lw.index_topk]``) and,
+        # with a mono dense prefix, the dense layers' (``self.k4[n]["dense_pre0"
+        # | "dense_pre" | "dense_post"]``), built and compiled on its first step
         self.k4 = WidthBuilds(
             self._build_layer_kernels, self._layer_kernels, group, "MiniMax-M3 mono"
         )
@@ -145,9 +164,10 @@ def _build_layer_kernels(self, n: int) -> dict:
             index_topk: self._build_k4(n, index_topk=index_topk)
             for index_topk in {lw.index_topk for lw in self.sparse}
         }
-        kernels["dense_pre0"] = self._build_dense_pre(n, plain_norm=True)
-        kernels["dense_pre"] = self._build_dense_pre(n)
-        kernels["dense_post"] = self._build_dense_post(n)
+        if self.dense_mono:
+            kernels["dense_pre0"] = self._build_dense_pre(n, plain_norm=True)
+            kernels["dense_pre"] = self._build_dense_pre(n)
+            kernels["dense_post"] = self._build_dense_post(n)
         return kernels
 
     @staticmethod
@@ -157,20 +177,18 @@ def _layer_kernels(kernels: dict) -> list:
             (launch, _KERNEL_ABI.get(kind, K4_ABI)) for kind, launch in kernels.items()
         ]
 
-    def _bind(self, causal_lm) -> None:
+    def _bind(self, causal_lm, drafter, require_aux: bool) -> None:
         """This rank's checks and weight views; raises ``MonoUnsupported``."""
         model = causal_lm.model
         config = causal_lm.config
         npes = get_tensor_model_parallel_world_size()
         if npes != TP:
             raise MonoUnsupported(f"TP {npes}")
         layers = model.layers
-        if len(layers) <= N_DENSE or any(
-            getattr(layers[i].self_attn, "is_indexed_sparse_attention", False)
-            or layers[i].is_moe_layer
-            for i in range(N_DENSE)
-        ):
-            raise MonoUnsupported("layers 0..2 must be dense full-attention layers")
+        self.dense_mono = dense_prefix_mono(layers)
+        if require_aux and drafter is None:
+            raise MonoUnsupported("no DSpark drafter")
+        self.aux = bind_aux_taps(drafter, config.hidden_size) if require_aux else None
         self.model = model
         self.npes = npes
         self.rank = get_tensor_model_parallel_rank()
@@ -188,7 +206,11 @@ def _bind(self, causal_lm) -> None:
         # (the original path raises on a cache miss there)
         if not self.sparse[0].index_topk:
             raise MonoUnsupported("the first sparse layer reuses a selection")
-        self.dense = [DenseLayer.from_layer(layers[i], i) for i in range(N_DENSE)]
+        self.dense = (
+            [DenseLayer.from_layer(layers[i], i) for i in range(N_DENSE)]
+            if self.dense_mono
+            else []
+        )
         for d in self.dense:
             # the decode dispatch reads what the original rope_cache step sets on
             # the impl; mono skips that step (a decode can precede every prefill:
@@ -199,7 +221,7 @@ def _bind(self, causal_lm) -> None:
                 or d.attn_impl.head_dim != 128
             )
         swiglu = {(d.swiglu_alpha, d.swiglu_beta, d.swiglu_limit) for d in self.dense}
-        if len(swiglu) != 1:
+        if self.dense_mono and len(swiglu) != 1:
             raise MonoUnsupported(f"dense layers' swiglu parameters {swiglu}")
         impl = self.sparse[0].attn_impl
         experts = layers[N_DENSE].block_sparse_moe.experts
@@ -246,22 +268,23 @@ def _bind(self, causal_lm) -> None:
             debug=self.debug,
         )
         # a dense layer's kernels around its original attention
-        self._build_dense_pre = functools.partial(
-            build_pre_attn_kernel,
-            eps,
-            self.sm_scale,
-            impl.init_blocks,
-            impl.local_blocks,
-            dense=True,
-        )
-        self._build_dense_post = functools.partial(
-            build_dense_post_kernel,
-            npes,
-            eps,
-            *swiglu.pop(),
-            debug=self.debug,
-            index_heads=heads.count,
-        )
+        if self.dense_mono:
+            self._build_dense_pre = functools.partial(
+                build_pre_attn_kernel,
+                eps,
+                self.sm_scale,
+                impl.init_blocks,
+                impl.local_blocks,
+                dense=True,
+            )
+            self._build_dense_post = functools.partial(
+                build_dense_post_kernel,
+                npes,
+                eps,
+                *swiglu.pop(),
+                debug=self.debug,
+                index_heads=heads.count,
+            )
 
     def _allocate(self, group) -> None:
         """Scratch, the peer handshake and the activation buffers (every rank)."""
@@ -327,10 +350,15 @@ def forward(
         aux: list[torch.Tensor] = []
         self.mailboxes.begin_step()
         fwd = get_forward_context()
-        res = self.run_dense_mono(input_ids, positions, fwd, aux_layers, aux)
+        res = (
+            self.run_dense_mono(input_ids, positions, fwd, aux_layers, aux)
+            if self.dense_mono
+            else self.run_dense(input_ids, positions, aux_layers, aux)
+        )
         rows = token_rows(fwd, res.shape[0])
         for i, lw in enumerate(self.sparse):
             res = self.run_sparse_layer(i, lw, fwd, positions, res, rows)
+            self._write_aux(lw.layer_id - 1, self.h[: res.shape[0]])
             if lw.layer_id in aux_layers:
                 aux.append(self.h[: res.shape[0]].clone())
         hidden = self.finish_step(res)
@@ -418,9 +446,9 @@ def run_dense(
         aux_layers: tuple[int, ...] = (),
         aux: list[torch.Tensor] | None = None,
     ) -> torch.Tensor:
-        """Embedding and the dense layers on the original modules; leaves the first
+        """Embedding and the prefix layers on the original modules; leaves the first
         sparse layer's (ar, res) in ``self.ars[0][:S]`` and the returned residual.
-        Appends the aux hidden state of the dense ones among ``aux_layers``."""
+        Appends the aux hidden state of the prefix ones among ``aux_layers``."""
         n = input_ids.numel()
         hidden = self.model.get_input_embeddings(input_ids)
         residual = None
@@ -478,11 +506,19 @@ def finish_step(self, res: torch.Tensor) -> torch.Tensor:
         """The final norm (the mailboxes are cleared at the next step's start:
         ``StepMailboxes.begin_step``)."""
         ar = self.ars[len(self.sparse) % 2]
-        hidden, _ = self.model.norm(ar[: res.shape[0]], res)
+        hidden, residual = self.model.norm(ar[: res.shape[0]], res)
+        self._write_aux(self.sparse[-1].layer_id, residual)
         if self.debug and not torch.cuda.is_current_stream_capturing():
             self.raise_on_given_up_waits(res.shape[0])
         return hidden
 
+    def _write_aux(self, layer_id: int, value: torch.Tensor) -> None:
+        if self.aux is None or layer_id not in self.aux.layer_ids:
+            return
+        self.aux.buffers[self.aux.layer_ids.index(layer_id)][: value.shape[0]].copy_(
+            value
+        )
+
     def raise_on_given_up_waits(self, n: int) -> None:
         """A debug build's step: raise with every mailbox wait that gave up (its
         region, pair index, tags, CTA) and every rank the fence gave up on."""

atom/models/registry.py

@@ -43,9 +43,8 @@
         "atom.models.glm5_next.Glm5NextForConditionalGeneration"
     ),
 }
-# Mono decode by wrapping the loaded model: the installer runs after the drafter
-# armed its hooks on the unwrapped layers. Architectures not listed here route
-# mono decode inside their own model class instead.
+# Model-specific Mono setup runs after the drafter arms the original layers.
+# An installer may wrap the model or finish binding routing already inside it.
 mono_decode_installers = {
     "DeepseekV4ForCausalLM": (
         "atom.models.deepseek_v4.mono.dispatch.install_mono_decode"
@@ -56,4 +55,10 @@
     "KimiK3ForConditionalGeneration": (
         "atom.models.kimi_k3.mono.dispatch.install_mono_decode"
     ),
+    "MiniMaxM3SparseForCausalLM": (
+        "atom.models.minimax_m3.mono.dispatch.install_mono_decode"
+    ),
+    "MiniMaxM3SparseForConditionalGeneration": (
+        "atom.models.minimax_m3.mono.dispatch.install_mono_decode"
+    ),
 }