Add MTP support for Step 3.5 Flash

mlx_lm/models/step3p5.py

@@ -2,7 +2,7 @@
 
 from dataclasses import dataclass
 from functools import partial
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Tuple
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -60,6 +60,7 @@ class ModelArgs(BaseModelArgs):
     norm_expert_weight: bool = True
     swiglu_limits: Optional[List[float]] = None
     swiglu_limits_shared: Optional[List[float]] = None
+    num_nextn_predict_layers: int = 0
     tie_word_embeddings: bool = False
 
 
@@ -339,16 +340,17 @@ def __init__(self, args: ModelArgs):
         self.norm = ZeroCenteredRMSNorm(args.hidden_size, eps=args.rms_norm_eps)
 
         self._swa_idx = next(
-            (i for i, l in enumerate(self.layers) if l.is_sliding), None
+            (i for i, layer in enumerate(self.layers) if layer.is_sliding), None
         )
         self._full_idx = next(
-            (i for i, l in enumerate(self.layers) if not l.is_sliding), None
+            (i for i, layer in enumerate(self.layers) if not layer.is_sliding), None
         )
 
     def __call__(
         self,
         x: mx.array,
         cache: Optional[List[Any]] = None,
+        return_prenorm: bool = False,
     ) -> mx.array:
         h = self.embed_tokens(x)
 
@@ -370,9 +372,92 @@ def __call__(
             mask = swa_mask if layer.is_sliding else full_mask
             h = layer(h, mask=mask, cache=c)
 
+        if return_prenorm:
+            return self.norm(h), h
         return self.norm(h)
 
 
+class Step3p5SharedHead(nn.Module):
+    """Per-MTP-layer prediction head with norm and output projection."""
+
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.norm = ZeroCenteredRMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.output = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.output(self.norm(x))
+
+
+class Step3p5MTPLayer(nn.Module):
+    """Single MTP prediction layer.
+
+    Architecture:
+      1. Normalize hidden_states (hnorm) and token embedding (enorm)
+      2. Concatenate and project: [B, L, 2H] -> [B, L, H] via eh_proj
+      3. Decoder block: sliding attention + dense MLP
+      4. Per-layer shared_head for logit prediction
+    """
+
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.hnorm = ZeroCenteredRMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.enorm = ZeroCenteredRMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.eh_proj = nn.Linear(args.hidden_size * 2, args.hidden_size, bias=False)
+
+        # MTP uses sliding_attention — pick first sliding layer_idx for RoPE config
+        mtp_layer_idx = 1
+        layer_types = args.layer_types or []
+        for idx, lt in enumerate(layer_types):
+            if lt == "sliding_attention":
+                mtp_layer_idx = idx
+                break
+
+        self.self_attn = Step3p5Attention(args, layer_idx=mtp_layer_idx)
+        self.mlp = Step3p5MLP(
+            args, intermediate_size=args.intermediate_size, swiglu_limit=0
+        )
+        self.input_layernorm = ZeroCenteredRMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+        self.post_attention_layernorm = ZeroCenteredRMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+        self.shared_head = Step3p5SharedHead(args)
+
+    def __call__(
+        self,
+        hidden_states: mx.array,
+        input_embeds: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> Tuple[mx.array, mx.array]:
+        h = self.hnorm(hidden_states)
+        e = self.enorm(input_embeds)
+        x = self.eh_proj(mx.concatenate([e, h], axis=-1))
+
+        residual = x
+        x = self.input_layernorm(x)
+        x = self.self_attn(x, mask=mask, cache=cache) + residual
+
+        residual = x
+        x = self.post_attention_layernorm(x)
+        x = self.mlp(x) + residual
+
+        logits = self.shared_head(x)
+        return x, logits
+
+
+class Step3p5MTP(nn.Module):
+    """MTP module with multiple prediction layers."""
+
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.layers = [
+            Step3p5MTPLayer(args) for _ in range(args.num_nextn_predict_layers)
+        ]
+
+
 class Model(nn.Module):
     def __init__(self, args: ModelArgs):
         super().__init__()
@@ -381,13 +466,30 @@ def __init__(self, args: ModelArgs):
         self.model = Step3p5Model(args)
         self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
 
+        self._mtp_num_layers = args.num_nextn_predict_layers
+        if self._mtp_num_layers > 0:
+            self.mtp = Step3p5MTP(args)
+        else:
+            self.mtp = None
+
     def __call__(
         self,
         inputs: mx.array,
         cache: Optional[List[Any]] = None,
+        return_hidden: bool = False,
     ):
-        out = self.model(inputs, cache)
-        return self.lm_head(out)
+        if return_hidden:
+            hidden_states, prenorm_hidden = self.model(
+                inputs, cache, return_prenorm=True
+            )
+        else:
+            hidden_states = self.model(inputs, cache)
+
+        out = self.lm_head(hidden_states)
+
+        if return_hidden:
+            return out, prenorm_hidden
+        return out
 
     @property
     def layers(self):
@@ -403,7 +505,54 @@ def make_cache(self):
             for layer in self.layers
         ]
 
+    def mtp_forward(
+        self,
+        hidden_states: mx.array,
+        next_token_ids: mx.array,
+        mtp_cache: Optional[Any] = None,
+    ) -> mx.array:
+        """Run MTP head to predict token n+2 given hidden states and token n+1.
+
+        Args:
+            hidden_states: [B, 1, H] prenorm hidden states from main model
+            next_token_ids: [B, 1] token IDs for position n+1
+            mtp_cache: list of KVCache for MTP layers
+
+        Returns:
+            logits: [B, 1, V] logits for token n+2
+        """
+        if self.mtp is None:
+            raise RuntimeError("MTP head not loaded (num_nextn_predict_layers=0)")
+
+        input_embeds = self.model.embed_tokens(next_token_ids)
+
+        layer = self.mtp.layers[0]
+        cache_entry = mtp_cache[0] if mtp_cache else None
+        mask = create_attention_mask(input_embeds, cache_entry)
+        _, logits = layer(hidden_states, input_embeds, mask=mask, cache=cache_entry)
+        return logits
+
+    def make_mtp_cache(self):
+        if self.mtp is None:
+            return None
+        # Mirror make_cache. An MTP layer borrows a real layer's attention
+        # config, so when that layer slides the MTP cache has to be windowed
+        # too: a plain KVCache would grow without bound across a generation
+        # while the main model's equivalent layers stay capped at
+        # sliding_window. Keyed off the attention module rather than assumed,
+        # because which layer it borrows depends on layer_types.
+        return [
+            (
+                RotatingKVCache(max_size=self.args.sliding_window)
+                if layer.self_attn.is_sliding
+                else KVCache()
+            )
+            for layer in self.mtp.layers
+        ]
+
     def sanitize(self, weights):
+        has_mtp_weights = any(k.startswith("mtp.") for k in weights)
+
         remappings = [
             (".moe.gate_proj.", ".mlp.switch_mlp.gate_proj."),
             (".moe.up_proj.", ".mlp.switch_mlp.up_proj."),
@@ -419,7 +568,8 @@ def sanitize(self, weights):
 
         new_weights = {}
         for k, v in weights.items():
-            if ".mtp" in k:
+            # Filter MTP weights if no MTP head or no MTP weights
+            if "mtp" in k and (not has_mtp_weights or self._mtp_num_layers == 0):
                 continue
             if "model.layers." in k:
                 parts = k.split(".")
@@ -451,6 +601,20 @@ def quant_predicate(self):
         def predicate(path, _):
             if "mlp.gate.gate" in path:
                 return {"group_size": 64, "bits": 8}
+            if "mtp." in path and any(
+                x in path
+                for x in [
+                    ".eh_proj.",
+                    ".enorm.",
+                    ".hnorm.",
+                    ".shared_head.norm.",
+                    ".input_layernorm.",
+                    ".post_attention_layernorm.",
+                    ".q_norm.",
+                    ".k_norm.",
+                ]
+            ):
+                return False
             return True
 
         return predicate

tests/test_models.py

@@ -1052,6 +1052,71 @@ def test_step3p5(self):
             model, args.model_type, args.vocab_size, args.num_hidden_layers
         )
 
+    def test_step3p5_mtp_cache_follows_the_main_model_policy(self):
+        """An MTP layer borrows a real layer's attention config.
+
+        When that layer slides, the MTP cache has to be windowed too, or it
+        grows without bound across a generation while the main model's
+        equivalent layers stay capped at sliding_window.
+        """
+        from mlx_lm.models import step3p5
+
+        def build(layer_types):
+            args = step3p5.ModelArgs(
+                model_type="step3p5",
+                hidden_size=256,
+                num_hidden_layers=4,
+                vocab_size=1024,
+                num_attention_heads=4,
+                num_attention_groups=2,
+                head_dim=64,
+                intermediate_size=512,
+                rms_norm_eps=1e-5,
+                rope_theta=[10000.0, 10000.0, 10000.0, 10000.0],
+                sliding_window=4,
+                layer_types=layer_types,
+                partial_rotary_factors=[0.5, 1.0, 1.0, 0.5],
+                attention_other_setting={
+                    "num_attention_heads": 8,
+                    "num_attention_groups": 2,
+                },
+                use_head_wise_attn_gate=True,
+                moe_num_experts=4,
+                moe_top_k=2,
+                moe_intermediate_size=256,
+                share_expert_dim=256,
+                moe_layers_enum="1,2,3",
+                num_nextn_predict_layers=1,
+            )
+            return step3p5.Model(args), args
+
+        sliding_model, args = build(
+            [
+                "full_attention",
+                "sliding_attention",
+                "sliding_attention",
+                "full_attention",
+            ]
+        )
+        self.assertTrue(sliding_model.mtp.layers[0].self_attn.is_sliding)
+        mtp_cache = sliding_model.make_mtp_cache()
+        self.assertIsInstance(mtp_cache[0], RotatingKVCache)
+        self.assertEqual(mtp_cache[0].max_size, args.sliding_window)
+
+        # It must actually stay bounded, not merely be the right class.
+        cache = mtp_cache[0]
+        for _ in range(4 * args.sliding_window):
+            cache.update_and_fetch(
+                mx.zeros((1, 2, 1, args.head_dim)), mx.zeros((1, 2, 1, args.head_dim))
+            )
+        self.assertLessEqual(cache.size(), args.sliding_window)
+
+        full_model, _ = build(["full_attention"] * 4)
+        self.assertFalse(full_model.mtp.layers[0].self_attn.is_sliding)
+        full_mtp_cache = full_model.make_mtp_cache()
+        self.assertIsInstance(full_mtp_cache[0], KVCache)
+        self.assertNotIsInstance(full_mtp_cache[0], RotatingKVCache)
+
     def test_step3p5_make_cache_uses_rotating_for_sliding_layers(self):
         from mlx_lm.models import step3p5