From 64153cabeeada666003b647db2c7b00edf71fc06 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?G=C3=B6kdeniz=20G=C3=BClmez?=
 <60228478+Goekdeniz-Guelmez@users.noreply.github.com>
Date: Mon, 16 Mar 2026 22:10:09 +0100
Subject: [PATCH 01/27] Add initial implementation of mistral4 model

---
 mlx_lm/models/mistral4.py | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 create mode 100644 mlx_lm/models/mistral4.py

diff --git a/mlx_lm/models/mistral4.py b/mlx_lm/models/mistral4.py
new file mode 100644
index 000000000..e69de29bb

From dfad789b62c87d83161a3cda3a8191a432f0c14e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?G=C3=B6kdeniz=20G=C3=BClmez?=
 <60228478+Goekdeniz-Guelmez@users.noreply.github.com>
Date: Mon, 16 Mar 2026 23:10:58 +0100
Subject: [PATCH 02/27] Implement Mistral-4 model architecture with attention
 and MoE layers

---
 mlx_lm/models/mistral4.py | 516 ++++++++++++++++++++++++++++++++++++++
 1 file changed, 516 insertions(+)

diff --git a/mlx_lm/models/mistral4.py b/mlx_lm/models/mistral4.py
index e69de29bb..0fc10d245 100644
--- a/mlx_lm/models/mistral4.py
+++ b/mlx_lm/models/mistral4.py
@@ -0,0 +1,516 @@
+# Copyright © 2026 Apple Inc.
+
+from dataclasses import dataclass
+import math
+from typing import Union, Dict, Optional, List, Any
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.nn.layers.distributed import shard_inplace, shard_linear
+
+from .base import BaseModelArgs, scaled_dot_product_attention
+from .rope_utils import initialize_rope
+from .deepseek_v3 import (
+    DeepseekV3MLP,
+    DeepseekV3Model,
+)
+from .switch_layers import SwitchGLU
+
+
+def _get_llama_4_attn_scale(size, offset, beta: float, max_position_embeddings: int):
+    if isinstance(offset, mx.array) and offset.ndim > 0:
+        offset = offset[:, None]
+
+    scaling = 1 + beta * mx.log(
+        1 + mx.floor((mx.arange(size) + offset) / max_position_embeddings)
+    )
+    if scaling.ndim == 2:
+        return scaling[:, None, :, None]
+    else:
+        return scaling[:, None]
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str = "mistral4"
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    moe_intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    n_shared_experts: int
+    n_routed_experts: int
+    routed_scaling_factor: float
+    kv_lora_rank: int
+    q_lora_rank: int
+    norm_topk_prob: bool
+    max_position_embeddings: int
+    rms_norm_eps: float
+    topk_group: int
+    num_experts_per_tok: int
+    first_k_dense_replace: int
+    n_group: int
+    qk_rope_head_dim: int
+    qk_nope_head_dim: int
+    v_head_dim: int
+    head_dim: Optional[int] = None
+    qk_head_dim: Optional[int] = None
+    rope_theta: float = 10000.0
+    tie_word_embeddings: bool = False
+    rope_parameters: Optional[Dict[str, Union[float, str, bool, List[int]]]] = None
+    rope_interleave: Optional[bool] = None
+    attention_bias: bool = False
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    rope_parameters: Optional[Dict[str, Union[float, str, bool, List[int]]]] = None
+
+    def __post_init__(self, **kwargs):
+        if self.num_key_value_groups is None:
+            self.num_key_value_groups = self.num_attention_heads // self.num_key_value_heads
+
+        if self.qk_head_dim is None:
+            self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+        
+        if self.head_dim is None:
+            self.head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+        
+        if self.rope_parameters is None:
+            self.rope_parameters = {
+                "type": "yarn",
+                "rope_theta": 10000.0,
+                "factor": 128.0,
+                "original_max_position_embeddings": 8192,
+                "max_position_embeddings": self.max_position_embeddings,
+                "beta_fast": 32.0,
+                "beta_slow": 1.0,
+                "mscale_all_dim": 1.0,
+                "mscale": 1.0,
+                "llama_4_scaling_beta": 0.1,
+                "partial_rotary_factor": self.qk_rope_head_dim / self.head_dim,
+            }
+
+
+class Mistral4Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.hidden_size = args.hidden_size
+        self.num_heads = args.num_attention_heads
+        self.max_position_embeddings = args.max_position_embeddings
+        self.rope_theta = args.rope_theta
+        self.q_lora_rank = args.q_lora_rank
+        self.qk_rope_head_dim = args.qk_rope_head_dim
+        self.kv_lora_rank = args.kv_lora_rank
+        self.v_head_dim = args.v_head_dim
+        self.qk_nope_head_dim = args.qk_nope_head_dim
+        self.q_head_dim = args.qk_nope_head_dim + args.qk_rope_head_dim
+
+        self.scale = self.q_head_dim**-0.5
+
+        if self.q_lora_rank is None:
+            self.q_proj = nn.Linear(
+                self.hidden_size, self.num_heads * self.q_head_dim, bias=False
+            )
+        else:
+            self.q_a_proj = nn.Linear(
+                self.hidden_size, self.q_lora_rank, bias=args.attention_bias
+            )
+            self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank, eps=1e-6)
+            self.q_b_proj = nn.Linear(
+                self.q_lora_rank, self.num_heads * self.q_head_dim, bias=False
+            )
+
+        self.kv_a_proj_with_mqa = nn.Linear(
+            self.hidden_size,
+            self.kv_lora_rank + self.qk_rope_head_dim,
+            bias=args.attention_bias,
+        )
+        self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank, eps=1e-6)
+        self.kv_b_proj = nn.Linear(
+            self.kv_lora_rank,
+            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
+            bias=False,
+        )
+
+        self.o_proj = nn.Linear(
+            self.num_heads * self.v_head_dim,
+            self.hidden_size,
+            bias=args.attention_bias,
+        )
+
+        if self.args.rope_parameters is not None:
+            mscale_all_dim = self.args.rope_parameters.get("mscale_all_dim", 0)
+            if mscale_all_dim:
+                scaling_factor = self.args.rope_parameters["factor"]
+                if scaling_factor > 1:
+                    s = 0.1 * mscale_all_dim * math.log(scaling_factor) + 1.0
+                    self.scale = self.scale * s * s
+
+        self.rope = initialize_rope(
+            dims=self.qk_rope_head_dim,
+            base=self.rope_theta,
+            traditional=True,
+            max_position_embeddings=self.max_position_embeddings,
+            scaling_args=self.args.rope_parameters,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        # Query projection
+        if self.q_lora_rank is None:
+            q = self.q_proj(x)
+        else:
+            q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x)))
+
+        q = q.reshape(B, L, self.num_heads, self.q_head_dim).transpose(0, 2, 1, 3)
+        q_nope, q_rope = mx.split(q, [self.qk_nope_head_dim], axis=-1)
+
+        # KV projection
+        compressed_kv = self.kv_a_proj_with_mqa(x)
+        k_latent, k_rope = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
+
+        # Project latent to K and V
+        kv = self.kv_b_proj(self.kv_a_layernorm(k_latent))
+        kv = kv.reshape(B, L, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
+        kv = kv.transpose(0, 2, 1, 3)
+        k_nope, v = mx.split(kv, [self.qk_nope_head_dim], axis=-1)
+
+        # Reshape k_rope to match k_nope shape
+        k_rope = k_rope.reshape(B, L, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3)
+
+        # Apply RoPE
+        offset = cache.offset if cache is not None else 0
+        q_rope = self.rope(q_rope, offset)
+        k_rope = self.rope(k_rope, offset)
+
+        # Expand k_rope to all heads
+        k_rope = mx.broadcast_to(k_rope, [B, self.num_heads, k_rope.shape[2], self.qk_rope_head_dim])
+
+        # Concatenate to form full query and key states
+        query_states = mx.concatenate([q_nope, q_rope], axis=-1)
+        key_states = mx.concatenate([k_nope, k_rope], axis=-1)
+
+        # Apply Llama-4 attention scaling
+        if self.args.rope_parameters is not None:
+            llama_4_beta = self.args.rope_parameters.get("llama_4_scaling_beta", 0.1)
+            original_max_pos = self.args.rope_parameters.get(
+                "original_max_position_embeddings", 8192
+            )
+            attn_scale = _get_llama_4_attn_scale(
+                L, offset, llama_4_beta, original_max_pos
+            )
+            query_states = query_states * attn_scale
+
+        # Update cache
+        if cache is not None:
+            key_states, v = cache.update_and_fetch(key_states, v)
+
+        # Standard attention
+        output = scaled_dot_product_attention(
+            query_states, key_states, v, scale=self.scale, mask=mask
+        )
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+@mx.compile
+def mistral4_expert_select(
+    gates,
+    top_k,
+    n_group,
+    topk_group,
+    routed_scaling_factor,
+    norm_topk_prob,
+):
+    """Mistral-4 routing using softmax instead of sigmoid."""
+    # Apply softmax to get normalized probabilities
+    scores = mx.softmax(gates.astype(mx.float32), axis=-1)
+    
+    if n_group > 1:
+        # Reshape to groups
+        scores_grouped = mx.unflatten(scores, axis=-1, shape=(n_group, -1))
+        # Get top-2 scores per group and sum them
+        group_scores = mx.topk(scores_grouped, 2, axis=-1).sum(axis=-1, keepdims=True)
+        
+        # Select top topk_group groups
+        group_idx = mx.topk(group_scores, topk_group, axis=-2)[1]
+        
+        # Create mask for selected groups
+        group_mask = mx.zeros_like(group_scores)
+        group_mask = mx.put_along_axis(
+            group_mask, group_idx, mx.array(1.0), axis=-2
+        )
+        
+        # Expand mask to expert dimension
+        score_mask = mx.flatten(
+            mx.broadcast_to(
+                group_mask,
+                group_scores.shape[:-1] + (n_group, scores_grouped.shape[-1])
+            ),
+            -2, -1
+        )
+        
+        # Mask out non-selected groups
+        scores_for_choice = scores * score_mask
+    else:
+        scores_for_choice = scores
+    
+    # Select top-k experts
+    inds = mx.topk(scores_for_choice, top_k, axis=-1)[1]
+    
+    # Gather weights from original scores
+    selected_scores = mx.take_along_axis(scores, inds, axis=-1)
+    
+    # Normalize if requested
+    if top_k > 1 and norm_topk_prob:
+        denominator = selected_scores.sum(axis=-1, keepdims=True) + 1e-20
+        selected_scores = selected_scores / denominator
+    
+    # Apply scaling factor
+    selected_scores = selected_scores * routed_scaling_factor
+    
+    return inds, selected_scores
+
+
+class Mistral4MoEGate(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.top_k = args.num_experts_per_tok
+        self.norm_topk_prob = args.norm_topk_prob
+        self.n_routed_experts = args.n_routed_experts
+        self.routed_scaling_factor = args.routed_scaling_factor
+        self.n_group = args.n_group
+        self.topk_group = args.topk_group
+        self.weight = mx.zeros((args.hidden_size, self.n_routed_experts))
+
+    def __call__(self, x):
+        gates = x @ self.weight
+        return mistral4_expert_select(
+            gates,
+            self.top_k,
+            self.n_group,
+            self.topk_group,
+            self.routed_scaling_factor,
+            self.norm_topk_prob,
+        )
+
+
+class Mistral4MoE(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.num_experts_per_tok = args.num_experts_per_tok
+        self.switch_mlp = SwitchGLU(
+            args.hidden_size,
+            args.moe_intermediate_size,
+            args.n_routed_experts,
+        )
+
+        self.gate = Mistral4MoEGate(args)
+        if args.n_shared_experts is not None:
+            intermediate_size = args.moe_intermediate_size * args.n_shared_experts
+            self.shared_experts = DeepseekV3MLP(
+                args=args, intermediate_size=intermediate_size
+            )
+
+    def __call__(self, x):
+        inds, scores = self.gate(x)
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
+        if self.args.n_shared_experts is not None:
+            y = y + self.shared_experts(x)
+        return y
+
+
+class Mistral4DecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.self_attn = Mistral4Attention(args)
+        self.mlp = (
+            Mistral4MoE(args)
+            if (
+                args.n_routed_experts is not None
+                and layer_idx >= args.first_k_dense_replace
+            )
+            else DeepseekV3MLP(args)
+        )
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        return h + r
+
+
+class Mistral4Model(DeepseekV3Model):
+    def __init__(self, args: ModelArgs):
+        nn.Module.__init__(self)
+        self.vocab_size = args.vocab_size
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            Mistral4DecoderLayer(args, idx)
+            for idx in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = Mistral4Model(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        out = self.model(inputs, cache=cache)
+        if self.args.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
+        return out
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def sanitize(self, weights):
+        def dequant(weight, scale_inv):
+            dtype = mx.bfloat16
+            weight = mx.from_fp8(weight, dtype=mx.bfloat16)
+            bs = 128
+            m, n = weight.shape
+            pad_bottom = (-m) % bs
+            pad_side = (-n) % bs
+            weight = mx.pad(weight, ((0, pad_bottom), (0, pad_side)))
+            weight = weight.reshape(
+                ((m + pad_bottom) // bs, bs, (n + pad_side) // bs, bs)
+            )
+            weight = (weight * scale_inv[:, None, :, None]).reshape(
+                m + pad_bottom, n + pad_side
+            )
+            return weight[:m, :n].astype(dtype)
+
+        # Remap for int4
+        new_weights = {}
+        for k, v in weights.items():
+            if k.endswith("weight_shape"):
+                base = k.replace("weight_shape", "")
+                new_weights[base + "weight"] = weights[base + "weight_packed"].view(
+                    mx.uint32
+                )
+                s = weights[base + "weight_scale"]
+                new_weights[base + "scales"] = s
+                new_weights[base + "biases"] = -8 * s
+            elif not (k.endswith("weight_scale") or k.endswith("weight_packed")):
+                new_weights[k] = v
+        weights = new_weights
+
+        # Dequantize fp8
+        new_weights = {}
+        for k, v in weights.items():
+            if "weight_scale_inv" in k:
+                scale_inv = v
+                wk = k.replace("_scale_inv", "")
+                weight = weights[wk]
+                weight = dequant(weight, scale_inv)
+                new_weights[wk] = weight
+            elif k not in new_weights:
+                new_weights[k] = v
+        weights = new_weights
+
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}"
+            for n, m in [("w1", "gate_proj"), ("w2", "down_proj"), ("w3", "up_proj")]:
+                for k in ["weight", "scales", "biases"]:
+                    if f"{prefix}.mlp.experts.0.{m}.{k}" in weights:
+                        to_join = [
+                            weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}")
+                            for e in range(self.args.n_routed_experts)
+                        ]
+                        weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
+
+        return {
+            k: v
+            for k, v in weights.items()
+            if "rotary_emb.inv_freq" not in k
+        }
+
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        rank = group.rank()
+        
+        for layer in self.model.layers:
+            if layer.self_attn.q_lora_rank is None:
+                layer.self_attn.q_proj = shard_linear(
+                    layer.self_attn.q_proj, "all-to-sharded", group=group
+                )
+            else:
+                layer.self_attn.q_b_proj = shard_linear(
+                    layer.self_attn.q_b_proj, "all-to-sharded", group=group
+                )
+            
+            layer.self_attn.kv_b_proj = shard_linear(
+                layer.self_attn.kv_b_proj, "all-to-sharded", group=group
+            )
+            
+            layer.self_attn.num_heads //= N
+
+            layer.self_attn.o_proj = shard_linear(
+                layer.self_attn.o_proj, "sharded-to-all", group=group
+            )
+
+            if isinstance(layer.mlp, DeepseekV3MLP):
+                layer.mlp.gate_proj = shard_linear(
+                    layer.mlp.gate_proj, "all-to-sharded", group=group
+                )
+                layer.mlp.down_proj = shard_linear(
+                    layer.mlp.down_proj, "sharded-to-all", group=group
+                )
+                layer.mlp.up_proj = shard_linear(
+                    layer.mlp.up_proj, "all-to-sharded", group=group
+                )
+
+            else:
+                if hasattr(layer.mlp, 'shared_experts'):
+                    shard_inplace(
+                        layer.mlp.shared_experts.gate_proj, "all-to-sharded", group=group
+                    )
+                    shard_inplace(
+                        layer.mlp.shared_experts.down_proj, "sharded-to-all", group=group
+                    )
+                    shard_inplace(
+                        layer.mlp.shared_experts.up_proj, "all-to-sharded", group=group
+                    )
+                shard_inplace(
+                    layer.mlp.switch_mlp.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.up_proj, "all-to-sharded", group=group
+                )
\ No newline at end of file

From 775547e0bf2c62cbf3d083603cef77ff4f46a03e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?G=C3=B6kdeniz=20G=C3=BClmez?=
 <60228478+Goekdeniz-Guelmez@users.noreply.github.com>
Date: Mon, 16 Mar 2026 23:12:37 +0100
Subject: [PATCH 03/27] Add Mistral AI's Mistral4 to acknowledgments in
 ACKNOWLEDGMENTS.md

---
 ACKNOWLEDGMENTS.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/ACKNOWLEDGMENTS.md b/ACKNOWLEDGMENTS.md
index 964053048..663387a2b 100644
--- a/ACKNOWLEDGMENTS.md
+++ b/ACKNOWLEDGMENTS.md
@@ -12,7 +12,7 @@ MLX LM was developed with contributions from the following individuals:
 OpenBMB's `MiniCPM` and `MiniCPM3`, Kyutai's `Helium`, State-Space's `Mamba v1` and 
 `Mamba v2`, Z.ai & THUKEG's `GLM`, `GLM4`, `GLM5 (GLM MoE DSA)`, Rednote `dots.llm1`, Baidu's `Ernie4.5 MoE`, 
 inclusionAI's `Bailing MoE e.g. Ling-family`, `Bailing MoE Linear e.g. Ling-Linear-family`, 
-Klear team - Kuaishou Technology's `Klear`, AI21 Lab's `Jamba` IBM's `Granite MoE`, 
+Klear team - Kuaishou Technology's `Klear`, AI21 Lab's `Jamba` IBM's `Granite MoE`, Mistral AI's `Mistral4`, 
 Meituan's `LongCat`, Nvidia's `Nemotron H`, Swiss-AI's `Apertus`, Nikity's `Lille130m`, 
 Alibaba Qwen's `Qwen3Next`, Tele-AI's `TeleChat3`, and Allenai's `OLMoE` and `Olmo 3`; 
 Helped add support for the following model architectures: 

From 5008bc0d7cf206cf1287e11de948687b826dcf43 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?G=C3=B6kdeniz=20G=C3=BClmez?=
 <60228478+Goekdeniz-Guelmez@users.noreply.github.com>
Date: Mon, 16 Mar 2026 23:33:28 +0100
Subject: [PATCH 04/27] Refactor Mistral4 model arguments and attention
 scaling; optimize expert selection logic

---
 mlx_lm/models/mistral4.py | 87 ++++++++++++++-------------------------
 1 file changed, 32 insertions(+), 55 deletions(-)

diff --git a/mlx_lm/models/mistral4.py b/mlx_lm/models/mistral4.py
index 0fc10d245..824ff52bc 100644
--- a/mlx_lm/models/mistral4.py
+++ b/mlx_lm/models/mistral4.py
@@ -1,7 +1,6 @@
 # Copyright © 2026 Apple Inc.
 
 from dataclasses import dataclass
-import math
 from typing import Union, Dict, Optional, List, Any
 
 import mlx.core as mx
@@ -31,7 +30,7 @@ def _get_llama_4_attn_scale(size, offset, beta: float, max_position_embeddings:
 
 @dataclass
 class ModelArgs(BaseModelArgs):
-    model_type: str = "mistral4"
+    model_type: str
     vocab_size: int
     hidden_size: int
     intermediate_size: int
@@ -58,16 +57,12 @@ class ModelArgs(BaseModelArgs):
     qk_head_dim: Optional[int] = None
     rope_theta: float = 10000.0
     tie_word_embeddings: bool = False
-    rope_parameters: Optional[Dict[str, Union[float, str, bool, List[int]]]] = None
     rope_interleave: Optional[bool] = None
     attention_bias: bool = False
     rope_scaling: Optional[Dict[str, Union[float, str]]] = None
     rope_parameters: Optional[Dict[str, Union[float, str, bool, List[int]]]] = None
 
     def __post_init__(self, **kwargs):
-        if self.num_key_value_groups is None:
-            self.num_key_value_groups = self.num_attention_heads // self.num_key_value_heads
-
         if self.qk_head_dim is None:
             self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
         
@@ -138,18 +133,10 @@ def __init__(self, args: ModelArgs):
             bias=args.attention_bias,
         )
 
-        if self.args.rope_parameters is not None:
-            mscale_all_dim = self.args.rope_parameters.get("mscale_all_dim", 0)
-            if mscale_all_dim:
-                scaling_factor = self.args.rope_parameters["factor"]
-                if scaling_factor > 1:
-                    s = 0.1 * mscale_all_dim * math.log(scaling_factor) + 1.0
-                    self.scale = self.scale * s * s
-
         self.rope = initialize_rope(
             dims=self.qk_rope_head_dim,
             base=self.rope_theta,
-            traditional=True,
+            traditional=args.rope_interleave if args.rope_interleave is not None else True,
             max_position_embeddings=self.max_position_embeddings,
             scaling_args=self.args.rope_parameters,
         )
@@ -229,53 +216,30 @@ def mistral4_expert_select(
     routed_scaling_factor,
     norm_topk_prob,
 ):
-    """Mistral-4 routing using softmax instead of sigmoid."""
-    # Apply softmax to get normalized probabilities
+    """Mistral-4 routing using softmax."""
     scores = mx.softmax(gates.astype(mx.float32), axis=-1)
-    
+
     if n_group > 1:
-        # Reshape to groups
         scores_grouped = mx.unflatten(scores, axis=-1, shape=(n_group, -1))
-        # Get top-2 scores per group and sum them
         group_scores = mx.topk(scores_grouped, 2, axis=-1).sum(axis=-1, keepdims=True)
-        
-        # Select top topk_group groups
-        group_idx = mx.topk(group_scores, topk_group, axis=-2)[1]
-        
-        # Create mask for selected groups
-        group_mask = mx.zeros_like(group_scores)
-        group_mask = mx.put_along_axis(
-            group_mask, group_idx, mx.array(1.0), axis=-2
+        # Zero out bottom (n_group - topk_group) groups
+        k = n_group - topk_group
+        group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :]
+        scores_grouped = mx.put_along_axis(
+            scores_grouped, mx.stop_gradient(group_idx), mx.array(0.0), axis=-2
         )
-        
-        # Expand mask to expert dimension
-        score_mask = mx.flatten(
-            mx.broadcast_to(
-                group_mask,
-                group_scores.shape[:-1] + (n_group, scores_grouped.shape[-1])
-            ),
-            -2, -1
-        )
-        
-        # Mask out non-selected groups
-        scores_for_choice = scores * score_mask
+        scores_for_choice = mx.flatten(scores_grouped, -2, -1)
     else:
         scores_for_choice = scores
-    
-    # Select top-k experts
-    inds = mx.topk(scores_for_choice, top_k, axis=-1)[1]
-    
-    # Gather weights from original scores
+
+    inds = mx.argpartition(-scores_for_choice, kth=top_k - 1, axis=-1)[..., :top_k]
+
     selected_scores = mx.take_along_axis(scores, inds, axis=-1)
-    
-    # Normalize if requested
-    if top_k > 1 and norm_topk_prob:
+    if norm_topk_prob:
         denominator = selected_scores.sum(axis=-1, keepdims=True) + 1e-20
         selected_scores = selected_scores / denominator
-    
-    # Apply scaling factor
     selected_scores = selected_scores * routed_scaling_factor
-    
+
     return inds, selected_scores
 
 
@@ -289,10 +253,10 @@ def __init__(self, args: ModelArgs):
         self.routed_scaling_factor = args.routed_scaling_factor
         self.n_group = args.n_group
         self.topk_group = args.topk_group
-        self.weight = mx.zeros((args.hidden_size, self.n_routed_experts))
+        self.weight = mx.zeros((self.n_routed_experts, args.hidden_size))
 
     def __call__(self, x):
-        gates = x @ self.weight
+        gates = x @ self.weight.T
         return mistral4_expert_select(
             gates,
             self.top_k,
@@ -318,7 +282,7 @@ def __init__(self, args: ModelArgs):
         if args.n_shared_experts is not None:
             intermediate_size = args.moe_intermediate_size * args.n_shared_experts
             self.shared_experts = DeepseekV3MLP(
-                args=args, intermediate_size=intermediate_size
+                args, intermediate_size=intermediate_size
             )
 
     def __call__(self, x):
@@ -443,7 +407,20 @@ def dequant(weight, scale_inv):
 
         for l in range(self.args.num_hidden_layers):
             prefix = f"model.layers.{l}"
-            for n, m in [("w1", "gate_proj"), ("w2", "down_proj"), ("w3", "up_proj")]:
+
+            # Handle fused gate_up_proj format (Mistral4NaiveMoe)
+            gup_key = f"{prefix}.mlp.experts.gate_up_proj"
+            if gup_key in weights:
+                gate_up = weights.pop(gup_key)
+                gate, up = mx.split(gate_up, 2, axis=1)
+                weights[f"{prefix}.mlp.switch_mlp.gate_proj.weight"] = gate
+                weights[f"{prefix}.mlp.switch_mlp.up_proj.weight"] = up
+            down_key = f"{prefix}.mlp.experts.down_proj"
+            if down_key in weights:
+                weights[f"{prefix}.mlp.switch_mlp.down_proj.weight"] = weights.pop(down_key)
+
+            # Handle per-expert weights format
+            for m in ["gate_proj", "down_proj", "up_proj"]:
                 for k in ["weight", "scales", "biases"]:
                     if f"{prefix}.mlp.experts.0.{m}.{k}" in weights:
                         to_join = [

From 1a3dac8062bbfff2b1a4251de3bf0009543b9405 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?G=C3=B6kdeniz=20G=C3=BClmez?=
 <60228478+Goekdeniz-Guelmez@users.noreply.github.com>
Date: Tue, 17 Mar 2026 13:42:55 +0100
Subject: [PATCH 05/27] Rename scaling_args to scaling_config in
 Mistral4Attention for clarity

---
 mlx_lm/models/mistral4.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlx_lm/models/mistral4.py b/mlx_lm/models/mistral4.py
index 824ff52bc..c490cfe2c 100644
--- a/mlx_lm/models/mistral4.py
+++ b/mlx_lm/models/mistral4.py
@@ -138,7 +138,7 @@ def __init__(self, args: ModelArgs):
             base=self.rope_theta,
             traditional=args.rope_interleave if args.rope_interleave is not None else True,
             max_position_embeddings=self.max_position_embeddings,
-            scaling_args=self.args.rope_parameters,
+            scaling_config=self.args.rope_parameters,
         )
 
     def __call__(

From d621dc0001c77cbd252dd873beafeecc7321f772 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?G=C3=B6kdeniz=20G=C3=BClmez?=
 <60228478+Goekdeniz-Guelmez@users.noreply.github.com>
Date: Tue, 17 Mar 2026 13:47:42 +0100
Subject: [PATCH 06/27] Integrate PipelineMixin into Mistral4Model and update
 attention call to include cache

---
 mlx_lm/models/mistral4.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/mlx_lm/models/mistral4.py b/mlx_lm/models/mistral4.py
index c490cfe2c..ea76ab2d3 100644
--- a/mlx_lm/models/mistral4.py
+++ b/mlx_lm/models/mistral4.py
@@ -14,6 +14,7 @@
     DeepseekV3Model,
 )
 from .switch_layers import SwitchGLU
+from .pipeline import PipelineMixin
 
 
 def _get_llama_4_attn_scale(size, offset, beta: float, max_position_embeddings: int):
@@ -200,7 +201,7 @@ def __call__(
 
         # Standard attention
         output = scaled_dot_product_attention(
-            query_states, key_states, v, scale=self.scale, mask=mask
+            query_states, key_states, v, cache=cache, scale=self.scale, mask=mask
         )
 
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
@@ -323,9 +324,9 @@ def __call__(
         return h + r
 
 
-class Mistral4Model(DeepseekV3Model):
+class Mistral4Model(DeepseekV3Model, PipelineMixin, nn.Module):
     def __init__(self, args: ModelArgs):
-        nn.Module.__init__(self)
+        PipelineMixin.__init__(self)
         self.vocab_size = args.vocab_size
         self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
         self.layers = [

From b4775790e493e14a96b847d56c1d894c85760af9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?G=C3=B6kdeniz=20G=C3=BClmez?=
 <60228478+Goekdeniz-Guelmez@users.noreply.github.com>
Date: Tue, 17 Mar 2026 13:53:06 +0100
Subject: [PATCH 07/27] Sanitize model weights by filtering and renaming keys
 in the sanitize method

---
 mlx_lm/models/mistral4.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/mlx_lm/models/mistral4.py b/mlx_lm/models/mistral4.py
index ea76ab2d3..b45cabd37 100644
--- a/mlx_lm/models/mistral4.py
+++ b/mlx_lm/models/mistral4.py
@@ -362,6 +362,17 @@ def layers(self):
         return self.model.layers
 
     def sanitize(self, weights):
+        sanitized = {}
+        for key, value in weights.items():
+            if key.startswith("vision_tower") or key.startswith("model.visual"):
+                continue
+            if key.startswith("model.language_model."):
+                key = key.replace("model.language_model.", "")
+            elif key.startswith("language_model."):
+                key = key.replace("language_model.", "")
+            sanitized[key] = value
+        weights = sanitized
+
         def dequant(weight, scale_inv):
             dtype = mx.bfloat16
             weight = mx.from_fp8(weight, dtype=mx.bfloat16)

From 55b0f20a69708704ddac85a61cbe746419337c2b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?G=C3=B6kdeniz=20G=C3=BClmez?=
 <60228478+Goekdeniz-Guelmez@users.noreply.github.com>
Date: Tue, 17 Mar 2026 13:53:19 +0100
Subject: [PATCH 08/27] format

---
 mlx_lm/models/mistral4.py | 50 ++++++++++++++++++++++-----------------
 1 file changed, 28 insertions(+), 22 deletions(-)

diff --git a/mlx_lm/models/mistral4.py b/mlx_lm/models/mistral4.py
index b45cabd37..f2126846c 100644
--- a/mlx_lm/models/mistral4.py
+++ b/mlx_lm/models/mistral4.py
@@ -1,20 +1,20 @@
 # Copyright © 2026 Apple Inc.
 
 from dataclasses import dataclass
-from typing import Union, Dict, Optional, List, Any
+from typing import Any, Dict, List, Optional, Union
 
 import mlx.core as mx
 import mlx.nn as nn
 from mlx.nn.layers.distributed import shard_inplace, shard_linear
 
 from .base import BaseModelArgs, scaled_dot_product_attention
-from .rope_utils import initialize_rope
 from .deepseek_v3 import (
     DeepseekV3MLP,
     DeepseekV3Model,
 )
-from .switch_layers import SwitchGLU
 from .pipeline import PipelineMixin
+from .rope_utils import initialize_rope
+from .switch_layers import SwitchGLU
 
 
 def _get_llama_4_attn_scale(size, offset, beta: float, max_position_embeddings: int):
@@ -29,6 +29,7 @@ def _get_llama_4_attn_scale(size, offset, beta: float, max_position_embeddings:
     else:
         return scaling[:, None]
 
+
 @dataclass
 class ModelArgs(BaseModelArgs):
     model_type: str
@@ -66,10 +67,10 @@ class ModelArgs(BaseModelArgs):
     def __post_init__(self, **kwargs):
         if self.qk_head_dim is None:
             self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
-        
+
         if self.head_dim is None:
             self.head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
-        
+
         if self.rope_parameters is None:
             self.rope_parameters = {
                 "type": "yarn",
@@ -137,7 +138,9 @@ def __init__(self, args: ModelArgs):
         self.rope = initialize_rope(
             dims=self.qk_rope_head_dim,
             base=self.rope_theta,
-            traditional=args.rope_interleave if args.rope_interleave is not None else True,
+            traditional=(
+                args.rope_interleave if args.rope_interleave is not None else True
+            ),
             max_position_embeddings=self.max_position_embeddings,
             scaling_config=self.args.rope_parameters,
         )
@@ -178,7 +181,9 @@ def __call__(
         k_rope = self.rope(k_rope, offset)
 
         # Expand k_rope to all heads
-        k_rope = mx.broadcast_to(k_rope, [B, self.num_heads, k_rope.shape[2], self.qk_rope_head_dim])
+        k_rope = mx.broadcast_to(
+            k_rope, [B, self.num_heads, k_rope.shape[2], self.qk_rope_head_dim]
+        )
 
         # Concatenate to form full query and key states
         query_states = mx.concatenate([q_nope, q_rope], axis=-1)
@@ -330,8 +335,7 @@ def __init__(self, args: ModelArgs):
         self.vocab_size = args.vocab_size
         self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
         self.layers = [
-            Mistral4DecoderLayer(args, idx)
-            for idx in range(args.num_hidden_layers)
+            Mistral4DecoderLayer(args, idx) for idx in range(args.num_hidden_layers)
         ]
         self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
 
@@ -429,7 +433,9 @@ def dequant(weight, scale_inv):
                 weights[f"{prefix}.mlp.switch_mlp.up_proj.weight"] = up
             down_key = f"{prefix}.mlp.experts.down_proj"
             if down_key in weights:
-                weights[f"{prefix}.mlp.switch_mlp.down_proj.weight"] = weights.pop(down_key)
+                weights[f"{prefix}.mlp.switch_mlp.down_proj.weight"] = weights.pop(
+                    down_key
+                )
 
             # Handle per-expert weights format
             for m in ["gate_proj", "down_proj", "up_proj"]:
@@ -441,17 +447,13 @@ def dequant(weight, scale_inv):
                         ]
                         weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
 
-        return {
-            k: v
-            for k, v in weights.items()
-            if "rotary_emb.inv_freq" not in k
-        }
+        return {k: v for k, v in weights.items() if "rotary_emb.inv_freq" not in k}
 
     def shard(self, group: Optional[mx.distributed.Group] = None):
         group = group or mx.distributed.init()
         N = group.size()
         rank = group.rank()
-        
+
         for layer in self.model.layers:
             if layer.self_attn.q_lora_rank is None:
                 layer.self_attn.q_proj = shard_linear(
@@ -461,11 +463,11 @@ def shard(self, group: Optional[mx.distributed.Group] = None):
                 layer.self_attn.q_b_proj = shard_linear(
                     layer.self_attn.q_b_proj, "all-to-sharded", group=group
                 )
-            
+
             layer.self_attn.kv_b_proj = shard_linear(
                 layer.self_attn.kv_b_proj, "all-to-sharded", group=group
             )
-            
+
             layer.self_attn.num_heads //= N
 
             layer.self_attn.o_proj = shard_linear(
@@ -484,12 +486,16 @@ def shard(self, group: Optional[mx.distributed.Group] = None):
                 )
 
             else:
-                if hasattr(layer.mlp, 'shared_experts'):
+                if hasattr(layer.mlp, "shared_experts"):
                     shard_inplace(
-                        layer.mlp.shared_experts.gate_proj, "all-to-sharded", group=group
+                        layer.mlp.shared_experts.gate_proj,
+                        "all-to-sharded",
+                        group=group,
                     )
                     shard_inplace(
-                        layer.mlp.shared_experts.down_proj, "sharded-to-all", group=group
+                        layer.mlp.shared_experts.down_proj,
+                        "sharded-to-all",
+                        group=group,
                     )
                     shard_inplace(
                         layer.mlp.shared_experts.up_proj, "all-to-sharded", group=group
@@ -502,4 +508,4 @@ def shard(self, group: Optional[mx.distributed.Group] = None):
                 )
                 shard_inplace(
                     layer.mlp.switch_mlp.up_proj, "all-to-sharded", group=group
-                )
\ No newline at end of file
+                )

From 352fbf50f88347f0a08e22845b24458bc8e5bee0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?G=C3=B6kdeniz=20G=C3=BClmez?=
 <60228478+Goekdeniz-Guelmez@users.noreply.github.com>
Date: Tue, 17 Mar 2026 13:58:15 +0100
Subject: [PATCH 09/27] Fix attention scaling in Mistral4Attention and add unit
 test for Mistral4 model

---
 mlx_lm/models/mistral4.py |  2 +-
 tests/test_models.py      | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/mlx_lm/models/mistral4.py b/mlx_lm/models/mistral4.py
index f2126846c..0aa511da6 100644
--- a/mlx_lm/models/mistral4.py
+++ b/mlx_lm/models/mistral4.py
@@ -198,7 +198,7 @@ def __call__(
             attn_scale = _get_llama_4_attn_scale(
                 L, offset, llama_4_beta, original_max_pos
             )
-            query_states = query_states * attn_scale
+            query_states = query_states * attn_scale.astype(query_states.dtype)
 
         # Update cache
         if cache is not None:
diff --git a/tests/test_models.py b/tests/test_models.py
index 0b2a963dc..e9a663db4 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -1192,6 +1192,39 @@ def test_deepseek_v3(self):
             model, args.model_type, args.vocab_size, args.num_hidden_layers
         )
 
+    def test_mistral4(self):
+        from mlx_lm.models import mistral4
+
+        args = mistral4.ModelArgs(
+            model_type="mistral4",
+            vocab_size=1024,
+            hidden_size=128,
+            intermediate_size=256,
+            moe_intermediate_size=256,
+            num_hidden_layers=4,
+            num_attention_heads=4,
+            num_key_value_heads=2,
+            n_routed_experts=4,
+            n_group=2,
+            topk_group=1,
+            num_experts_per_tok=2,
+            n_shared_experts=1,
+            routed_scaling_factor=1.0,
+            kv_lora_rank=4,
+            q_lora_rank=4,
+            qk_rope_head_dim=32,
+            v_head_dim=16,
+            qk_nope_head_dim=32,
+            norm_topk_prob=True,
+            max_position_embeddings=4096,
+            rms_norm_eps=1e-6,
+            first_k_dense_replace=0,
+        )
+        model = mistral4.Model(args)
+        self.model_test_runner(
+            model, args.model_type, args.vocab_size, args.num_hidden_layers
+        )
+
     def test_gemma2(self):
         from mlx_lm.models import gemma2
 

From 79082fa78c9e3e307e9aa8239b7b6006efd8e0f1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?G=C3=B6kdeniz=20G=C3=BClmez?=
 <60228478+Goekdeniz-Guelmez@users.noreply.github.com>
Date: Tue, 17 Mar 2026 13:59:39 +0100
Subject: [PATCH 10/27] Refactor Mistral4Attention by removing commented-out
 code and unnecessary comments for clarity

---
 mlx_lm/models/mistral4.py | 8 --------
 1 file changed, 8 deletions(-)

diff --git a/mlx_lm/models/mistral4.py b/mlx_lm/models/mistral4.py
index 0aa511da6..72a7c0431 100644
--- a/mlx_lm/models/mistral4.py
+++ b/mlx_lm/models/mistral4.py
@@ -153,7 +153,6 @@ def __call__(
     ) -> mx.array:
         B, L, D = x.shape
 
-        # Query projection
         if self.q_lora_rank is None:
             q = self.q_proj(x)
         else:
@@ -162,7 +161,6 @@ def __call__(
         q = q.reshape(B, L, self.num_heads, self.q_head_dim).transpose(0, 2, 1, 3)
         q_nope, q_rope = mx.split(q, [self.qk_nope_head_dim], axis=-1)
 
-        # KV projection
         compressed_kv = self.kv_a_proj_with_mqa(x)
         k_latent, k_rope = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
 
@@ -172,10 +170,8 @@ def __call__(
         kv = kv.transpose(0, 2, 1, 3)
         k_nope, v = mx.split(kv, [self.qk_nope_head_dim], axis=-1)
 
-        # Reshape k_rope to match k_nope shape
         k_rope = k_rope.reshape(B, L, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3)
 
-        # Apply RoPE
         offset = cache.offset if cache is not None else 0
         q_rope = self.rope(q_rope, offset)
         k_rope = self.rope(k_rope, offset)
@@ -185,7 +181,6 @@ def __call__(
             k_rope, [B, self.num_heads, k_rope.shape[2], self.qk_rope_head_dim]
         )
 
-        # Concatenate to form full query and key states
         query_states = mx.concatenate([q_nope, q_rope], axis=-1)
         key_states = mx.concatenate([k_nope, k_rope], axis=-1)
 
@@ -200,11 +195,9 @@ def __call__(
             )
             query_states = query_states * attn_scale.astype(query_states.dtype)
 
-        # Update cache
         if cache is not None:
             key_states, v = cache.update_and_fetch(key_states, v)
 
-        # Standard attention
         output = scaled_dot_product_attention(
             query_states, key_states, v, cache=cache, scale=self.scale, mask=mask
         )
@@ -222,7 +215,6 @@ def mistral4_expert_select(
     routed_scaling_factor,
     norm_topk_prob,
 ):
-    """Mistral-4 routing using softmax."""
     scores = mx.softmax(gates.astype(mx.float32), axis=-1)
 
     if n_group > 1:

From 443a1bc3788585214f6b1f9cab5273ad66471860 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?G=C3=B6kdeniz=20G=C3=BClmez?=
 <60228478+Goekdeniz-Guelmez@users.noreply.github.com>
Date: Wed, 18 Mar 2026 15:24:06 +0100
Subject: [PATCH 11/27] Add Mistral4 model implementation and integrate with
 existing architecture

- Introduced Mistral4 model in a new file `mistral4_text.py`.
- Updated model imports in `mistral3.py` to include Mistral4.
- Modified `utils.py` to remap "mistral4" to "mistral3".
- Updated test cases in `test_models.py` to test Mistral4 using the new model class.
- Removed the old Mistral4 implementation from `mistral4.py`.
- Added attention scaling and model architecture specific to Mistral4.
---
 mlx_lm/generate.py                            |  1 +
 mlx_lm/models/mistral3.py                     |  6 +++-
 .../models/{mistral4.py => mistral4_text.py}  | 33 +++++++++++--------
 mlx_lm/utils.py                               |  1 +
 tests/test_models.py                          |  6 ++--
 5 files changed, 30 insertions(+), 17 deletions(-)
 rename mlx_lm/models/{mistral4.py => mistral4_text.py} (95%)

diff --git a/mlx_lm/generate.py b/mlx_lm/generate.py
index ef8dbf7bf..9309671be 100644
--- a/mlx_lm/generate.py
+++ b/mlx_lm/generate.py
@@ -1455,6 +1455,7 @@ def main():
         adapter_path=args.adapter_path,
         tokenizer_config=tokenizer_config,
         model_config={"quantize_activations": args.quantize_activations},
+        lazy=True,
     )
     for eos_token in args.extra_eos_token:
         tokenizer.add_eos_token(eos_token)
diff --git a/mlx_lm/models/mistral3.py b/mlx_lm/models/mistral3.py
index b2d93543e..422065767 100644
--- a/mlx_lm/models/mistral3.py
+++ b/mlx_lm/models/mistral3.py
@@ -7,7 +7,7 @@
 import mlx.nn as nn
 from mlx.utils import tree_flatten, tree_unflatten
 
-from . import llama, ministral3
+from . import llama, ministral3, mistral4_text
 from .base import BaseModelArgs
 
 
@@ -30,6 +30,10 @@ def __init__(self, args: ModelArgs):
             self.language_model = ministral3.Model(
                 ministral3.ModelArgs.from_dict(args.text_config)
             )
+        elif args.text_config.get("model_type") == "mistral4":
+            self.language_model = mistral4_text.Model(
+                mistral4_text.ModelArgs.from_dict(args.text_config)
+            )
         else:
             self.language_model = llama.Model(
                 llama.ModelArgs.from_dict(args.text_config)
diff --git a/mlx_lm/models/mistral4.py b/mlx_lm/models/mistral4_text.py
similarity index 95%
rename from mlx_lm/models/mistral4.py
rename to mlx_lm/models/mistral4_text.py
index 72a7c0431..252d38c40 100644
--- a/mlx_lm/models/mistral4.py
+++ b/mlx_lm/models/mistral4_text.py
@@ -7,7 +7,7 @@
 import mlx.nn as nn
 from mlx.nn.layers.distributed import shard_inplace, shard_linear
 
-from .base import BaseModelArgs, scaled_dot_product_attention
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .deepseek_v3 import (
     DeepseekV3MLP,
     DeepseekV3Model,
@@ -331,6 +331,23 @@ def __init__(self, args: ModelArgs):
         ]
         self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
 
+    def __call__(
+        self,
+        x: mx.array,
+        cache: Optional[Any] = None,
+        input_embeddings: Optional[mx.array] = None,
+    ) -> mx.array:
+        h = input_embeddings if input_embeddings is not None else self.embed_tokens(x)
+
+        if cache is None:
+            cache = [None] * len(self.pipeline_layers)
+        mask = create_attention_mask(h, cache[0], return_array=True)
+
+        for l, c in zip(self.pipeline_layers, cache):
+            h = l(h, mask, cache=c)
+
+        return self.norm(h)
+
 
 class Model(nn.Module):
     def __init__(self, args: ModelArgs):
@@ -345,8 +362,9 @@ def __call__(
         self,
         inputs: mx.array,
         cache: Optional[Any] = None,
+        input_embeddings: Optional[mx.array] = None,
     ) -> mx.array:
-        out = self.model(inputs, cache=cache)
+        out = self.model(inputs, cache=cache, input_embeddings=input_embeddings)
         if self.args.tie_word_embeddings:
             out = self.model.embed_tokens.as_linear(out)
         else:
@@ -358,17 +376,6 @@ def layers(self):
         return self.model.layers
 
     def sanitize(self, weights):
-        sanitized = {}
-        for key, value in weights.items():
-            if key.startswith("vision_tower") or key.startswith("model.visual"):
-                continue
-            if key.startswith("model.language_model."):
-                key = key.replace("model.language_model.", "")
-            elif key.startswith("language_model."):
-                key = key.replace("language_model.", "")
-            sanitized[key] = value
-        weights = sanitized
-
         def dequant(weight, scale_inv):
             dtype = mx.bfloat16
             weight = mx.from_fp8(weight, dtype=mx.bfloat16)
diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py
index 70bf8c83f..ffac0eccb 100644
--- a/mlx_lm/utils.py
+++ b/mlx_lm/utils.py
@@ -44,6 +44,7 @@
 # Constants
 MODEL_REMAPPING = {
     "mistral": "llama",
+    "mistral4": "mistral3",
     "llava": "mistral3",
     "phi-msft": "phixtral",
     "falcon_mamba": "mamba",
diff --git a/tests/test_models.py b/tests/test_models.py
index e9a663db4..25283cd06 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -1193,9 +1193,9 @@ def test_deepseek_v3(self):
         )
 
     def test_mistral4(self):
-        from mlx_lm.models import mistral4
+        from mlx_lm.models import mistral4_text
 
-        args = mistral4.ModelArgs(
+        args = mistral4_text.ModelArgs(
             model_type="mistral4",
             vocab_size=1024,
             hidden_size=128,
@@ -1220,7 +1220,7 @@ def test_mistral4(self):
             rms_norm_eps=1e-6,
             first_k_dense_replace=0,
         )
-        model = mistral4.Model(args)
+        model = mistral4_text.Model(args)
         self.model_test_runner(
             model, args.model_type, args.vocab_size, args.num_hidden_layers
         )

From ec42929c2b37c6a392c53bbdee41fba154637a0f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?G=C3=B6kdeniz=20G=C3=BClmez?=
 <60228478+Goekdeniz-Guelmez@users.noreply.github.com>
Date: Wed, 18 Mar 2026 15:33:35 +0100
Subject: [PATCH 12/27] Update Mistral4 model references in utils and remove
 unused function from mistral4_text

---
 mlx_lm/models/mistral4_text.py | 14 +-------------
 mlx_lm/utils.py                |  2 +-
 2 files changed, 2 insertions(+), 14 deletions(-)

diff --git a/mlx_lm/models/mistral4_text.py b/mlx_lm/models/mistral4_text.py
index 252d38c40..613206747 100644
--- a/mlx_lm/models/mistral4_text.py
+++ b/mlx_lm/models/mistral4_text.py
@@ -12,24 +12,12 @@
     DeepseekV3MLP,
     DeepseekV3Model,
 )
+from .ministral3 import _get_llama_4_attn_scale
 from .pipeline import PipelineMixin
 from .rope_utils import initialize_rope
 from .switch_layers import SwitchGLU
 
 
-def _get_llama_4_attn_scale(size, offset, beta: float, max_position_embeddings: int):
-    if isinstance(offset, mx.array) and offset.ndim > 0:
-        offset = offset[:, None]
-
-    scaling = 1 + beta * mx.log(
-        1 + mx.floor((mx.arange(size) + offset) / max_position_embeddings)
-    )
-    if scaling.ndim == 2:
-        return scaling[:, None, :, None]
-    else:
-        return scaling[:, None]
-
-
 @dataclass
 class ModelArgs(BaseModelArgs):
     model_type: str
diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py
index ffac0eccb..44c482896 100644
--- a/mlx_lm/utils.py
+++ b/mlx_lm/utils.py
@@ -44,7 +44,7 @@
 # Constants
 MODEL_REMAPPING = {
     "mistral": "llama",
-    "mistral4": "mistral3",
+    "mistral4": "mistral4_text",
     "llava": "mistral3",
     "phi-msft": "phixtral",
     "falcon_mamba": "mamba",

From 8d275113d02b8b8b3993507d7a47ff623b9c4aaf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?G=C3=B6kdeniz=20G=C3=BClmez?=
 <60228478+Goekdeniz-Guelmez@users.noreply.github.com>
Date: Wed, 18 Mar 2026 15:37:56 +0100
Subject: [PATCH 13/27] Enhance Mistral4 attention mechanism by integrating
 attention scaling and updating method signatures

---
 mlx_lm/models/mistral4_text.py | 27 +++++++++++++++------------
 1 file changed, 15 insertions(+), 12 deletions(-)

diff --git a/mlx_lm/models/mistral4_text.py b/mlx_lm/models/mistral4_text.py
index 613206747..48e7189e3 100644
--- a/mlx_lm/models/mistral4_text.py
+++ b/mlx_lm/models/mistral4_text.py
@@ -136,6 +136,7 @@ def __init__(self, args: ModelArgs):
     def __call__(
         self,
         x: mx.array,
+        attn_scale: mx.array,
         mask: Optional[mx.array] = None,
         cache: Optional[Any] = None,
     ) -> mx.array:
@@ -172,16 +173,7 @@ def __call__(
         query_states = mx.concatenate([q_nope, q_rope], axis=-1)
         key_states = mx.concatenate([k_nope, k_rope], axis=-1)
 
-        # Apply Llama-4 attention scaling
-        if self.args.rope_parameters is not None:
-            llama_4_beta = self.args.rope_parameters.get("llama_4_scaling_beta", 0.1)
-            original_max_pos = self.args.rope_parameters.get(
-                "original_max_position_embeddings", 8192
-            )
-            attn_scale = _get_llama_4_attn_scale(
-                L, offset, llama_4_beta, original_max_pos
-            )
-            query_states = query_states * attn_scale.astype(query_states.dtype)
+        query_states = query_states * attn_scale
 
         if cache is not None:
             key_states, v = cache.update_and_fetch(key_states, v)
@@ -300,10 +292,11 @@ def __init__(self, args: ModelArgs, layer_idx: int):
     def __call__(
         self,
         x: mx.array,
+        attn_scale: mx.array,
         mask: Optional[mx.array] = None,
         cache: Optional[Any] = None,
     ) -> mx.array:
-        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        r = self.self_attn(self.input_layernorm(x), attn_scale, mask, cache)
         h = x + r
         r = self.mlp(self.post_attention_layernorm(h))
         return h + r
@@ -312,6 +305,7 @@ def __call__(
 class Mistral4Model(DeepseekV3Model, PipelineMixin, nn.Module):
     def __init__(self, args: ModelArgs):
         PipelineMixin.__init__(self)
+        self.args = args
         self.vocab_size = args.vocab_size
         self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
         self.layers = [
@@ -329,10 +323,19 @@ def __call__(
 
         if cache is None:
             cache = [None] * len(self.pipeline_layers)
+
+        offset = cache[0].offset if cache[0] is not None else 0
         mask = create_attention_mask(h, cache[0], return_array=True)
 
+        attn_scale = _get_llama_4_attn_scale(
+            x.shape[1],
+            offset,
+            self.args.rope_parameters["llama_4_scaling_beta"],
+            self.args.rope_parameters["original_max_position_embeddings"],
+        ).astype(h.dtype)
+
         for l, c in zip(self.pipeline_layers, cache):
-            h = l(h, mask, cache=c)
+            h = l(h, attn_scale, mask, cache=c)
 
         return self.norm(h)
 

From 80bd71338213df309d409efb17f6c085c776c451 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?G=C3=B6kdeniz=20G=C3=BClmez?=
 <60228478+Goekdeniz-Guelmez@users.noreply.github.com>
Date: Wed, 18 Mar 2026 15:41:08 +0100
Subject: [PATCH 14/27] Refactor ModelArgs to simplify rope_parameters handling
 and update default values

---
 mlx_lm/models/mistral4_text.py | 19 ++++---------------
 1 file changed, 4 insertions(+), 15 deletions(-)

diff --git a/mlx_lm/models/mistral4_text.py b/mlx_lm/models/mistral4_text.py
index 48e7189e3..f04108dae 100644
--- a/mlx_lm/models/mistral4_text.py
+++ b/mlx_lm/models/mistral4_text.py
@@ -50,7 +50,7 @@ class ModelArgs(BaseModelArgs):
     rope_interleave: Optional[bool] = None
     attention_bias: bool = False
     rope_scaling: Optional[Dict[str, Union[float, str]]] = None
-    rope_parameters: Optional[Dict[str, Union[float, str, bool, List[int]]]] = None
+    rope_parameters: Optional[Dict] = None
 
     def __post_init__(self, **kwargs):
         if self.qk_head_dim is None:
@@ -59,20 +59,9 @@ def __post_init__(self, **kwargs):
         if self.head_dim is None:
             self.head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
 
-        if self.rope_parameters is None:
-            self.rope_parameters = {
-                "type": "yarn",
-                "rope_theta": 10000.0,
-                "factor": 128.0,
-                "original_max_position_embeddings": 8192,
-                "max_position_embeddings": self.max_position_embeddings,
-                "beta_fast": 32.0,
-                "beta_slow": 1.0,
-                "mscale_all_dim": 1.0,
-                "mscale": 1.0,
-                "llama_4_scaling_beta": 0.1,
-                "partial_rotary_factor": self.qk_rope_head_dim / self.head_dim,
-            }
+        if self.rope_parameters is not None:
+            self.rope_theta = self.rope_parameters.get("rope_theta", 100000.0)
+            self.rope_scaling = self.rope_parameters
 
 
 class Mistral4Attention(nn.Module):

From 85d7b7a02feab7929e678bb49d9b513bfb85e9d3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?G=C3=B6kdeniz=20G=C3=BClmez?=
 <60228478+Goekdeniz-Guelmez@users.noreply.github.com>
Date: Thu, 19 Mar 2026 10:52:52 +0100
Subject: [PATCH 15/27] Refactor sanitize method to improve weight handling and
 remove unused utility functions

---
 mlx_lm/models/mistral3.py | 23 ++++++++++++++++-------
 1 file changed, 16 insertions(+), 7 deletions(-)

diff --git a/mlx_lm/models/mistral3.py b/mlx_lm/models/mistral3.py
index 422065767..cea3cdff5 100644
--- a/mlx_lm/models/mistral3.py
+++ b/mlx_lm/models/mistral3.py
@@ -5,7 +5,6 @@
 
 import mlx.core as mx
 import mlx.nn as nn
-from mlx.utils import tree_flatten, tree_unflatten
 
 from . import llama, ministral3, mistral4_text
 from .base import BaseModelArgs
@@ -50,12 +49,22 @@ def __call__(
         )
 
     def sanitize(self, weights):
-        weights = tree_unflatten(list(weights.items()))
-        weights.pop("vision_tower", None)
-        weights.pop("multi_modal_projector", None)
-        lm_weights = dict(tree_flatten(weights["language_model"]))
-        weights["language_model"] = self.language_model.sanitize(lm_weights)
-        return dict(tree_flatten(weights))
+        sanitized = {}
+        for key, value in weights.items():
+            if "vision_tower" in key or "multi_modal_projector" in key:
+                continue
+            if key.startswith("model."):
+                key = key[len("model."):]
+            sanitized[key] = value
+
+        lm_weights = {
+            k[len("language_model."):]: v
+            for k, v in sanitized.items()
+            if k.startswith("language_model.")
+        }
+
+        sanitized_lm = self.language_model.sanitize(lm_weights)
+        return {"language_model." + k: v for k, v in sanitized_lm.items()}
 
     @property
     def layers(self):

From 88e98c576a9e96f61f64148b55fd4bf42c1ef90f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?G=C3=B6kdeniz=20G=C3=BClmez?=
 <60228478+Goekdeniz-Guelmez@users.noreply.github.com>
Date: Thu, 19 Mar 2026 11:03:28 +0100
Subject: [PATCH 16/27] nits + format

---
 mlx_lm/models/mistral3.py      | 6 +++---
 mlx_lm/models/mistral4_text.py | 4 +---
 2 files changed, 4 insertions(+), 6 deletions(-)

diff --git a/mlx_lm/models/mistral3.py b/mlx_lm/models/mistral3.py
index cea3cdff5..ba986990d 100644
--- a/mlx_lm/models/mistral3.py
+++ b/mlx_lm/models/mistral3.py
@@ -1,4 +1,4 @@
-# Copyright © 2025 Apple Inc.
+# Copyright © 2026 Apple Inc.
 
 from dataclasses import dataclass
 from typing import Optional
@@ -54,11 +54,11 @@ def sanitize(self, weights):
             if "vision_tower" in key or "multi_modal_projector" in key:
                 continue
             if key.startswith("model."):
-                key = key[len("model."):]
+                key = key[len("model.") :]
             sanitized[key] = value
 
         lm_weights = {
-            k[len("language_model."):]: v
+            k[len("language_model.") :]: v
             for k, v in sanitized.items()
             if k.startswith("language_model.")
         }
diff --git a/mlx_lm/models/mistral4_text.py b/mlx_lm/models/mistral4_text.py
index f04108dae..9b65770e9 100644
--- a/mlx_lm/models/mistral4_text.py
+++ b/mlx_lm/models/mistral4_text.py
@@ -1,7 +1,7 @@
 # Copyright © 2026 Apple Inc.
 
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, Optional, Union
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -416,7 +416,6 @@ def dequant(weight, scale_inv):
                     down_key
                 )
 
-            # Handle per-expert weights format
             for m in ["gate_proj", "down_proj", "up_proj"]:
                 for k in ["weight", "scales", "biases"]:
                     if f"{prefix}.mlp.experts.0.{m}.{k}" in weights:
@@ -431,7 +430,6 @@ def dequant(weight, scale_inv):
     def shard(self, group: Optional[mx.distributed.Group] = None):
         group = group or mx.distributed.init()
         N = group.size()
-        rank = group.rank()
 
         for layer in self.model.layers:
             if layer.self_attn.q_lora_rank is None:

From 9ec6bea1dfd05c5907adc3e408ff6900a938639f Mon Sep 17 00:00:00 2001
From: Michal Klein <46717574+michalk8@users.noreply.github.com>
Date: Wed, 16 Sep 2026 15:36:21 +0200
Subject: [PATCH 17/27] Fix yarn for Mistral v4

---
 mlx_lm/models/mistral4_text.py |  4 ++--
 mlx_lm/models/rope_utils.py    | 17 +++++++++++++++++
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/mlx_lm/models/mistral4_text.py b/mlx_lm/models/mistral4_text.py
index 9b65770e9..58729abca 100644
--- a/mlx_lm/models/mistral4_text.py
+++ b/mlx_lm/models/mistral4_text.py
@@ -14,7 +14,7 @@
 )
 from .ministral3 import _get_llama_4_attn_scale
 from .pipeline import PipelineMixin
-from .rope_utils import initialize_rope
+from .rope_utils import apply_yarn_mscale, initialize_rope
 from .switch_layers import SwitchGLU
 
 
@@ -79,7 +79,7 @@ def __init__(self, args: ModelArgs):
         self.qk_nope_head_dim = args.qk_nope_head_dim
         self.q_head_dim = args.qk_nope_head_dim + args.qk_rope_head_dim
 
-        self.scale = self.q_head_dim**-0.5
+        self.scale = apply_yarn_mscale(self.q_head_dim**-0.5, args.rope_parameters)
 
         if self.q_lora_rank is None:
             self.q_proj = nn.Linear(
diff --git a/mlx_lm/models/rope_utils.py b/mlx_lm/models/rope_utils.py
index 04c8d08f8..d87bdb5c9 100644
--- a/mlx_lm/models/rope_utils.py
+++ b/mlx_lm/models/rope_utils.py
@@ -269,6 +269,23 @@ def __call__(self, x: mx.array, offset: int = 0) -> mx.array:
         )
 
 
+def apply_yarn_mscale(scale: float, scaling_config: Optional[dict]) -> float:
+    """Fold the yarn mscale into an attention scale.
+
+    ``initialize_rope`` puts the ``mscale / mscale_all_dim`` ratio on the rope;
+    this is the other half, ``mscale_all_dim`` squared.
+    """
+    if not scaling_config:
+        return scale
+    rope_type = scaling_config.get("type") or scaling_config.get("rope_type", "default")
+    mscale_all_dim = scaling_config.get("mscale_all_dim", 0)
+    factor = scaling_config.get("factor", 1)
+    if rope_type == "default" or not mscale_all_dim or factor <= 1:
+        return scale
+    s = 0.1 * mscale_all_dim * math.log(factor) + 1.0
+    return scale * s * s
+
+
 def initialize_rope(
     dims,
     base,

From fa201e85d15991f92ac43cbfdbe9902e75e87bc1 Mon Sep 17 00:00:00 2001
From: Michal Klein <46717574+michalk8@users.noreply.github.com>
Date: Thu, 17 Sep 2026 11:16:47 +0200
Subject: [PATCH 18/27] Fix Mistral4 sanitize

---
 mlx_lm/models/mistral4_text.py | 23 +++++++++++++++++------
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/mlx_lm/models/mistral4_text.py b/mlx_lm/models/mistral4_text.py
index 58729abca..7d39fb690 100644
--- a/mlx_lm/models/mistral4_text.py
+++ b/mlx_lm/models/mistral4_text.py
@@ -356,9 +356,18 @@ def layers(self):
         return self.model.layers
 
     def sanitize(self, weights):
+        def broadcasts(scale_shape, weight_shape):
+            if len(scale_shape) > len(weight_shape):
+                return False
+            pad = (1,) * (len(weight_shape) - len(scale_shape)) + tuple(scale_shape)
+            return all(s in (1, w) for s, w in zip(pad, weight_shape))
+
         def dequant(weight, scale_inv):
             dtype = mx.bfloat16
-            weight = mx.from_fp8(weight, dtype=mx.bfloat16)
+            weight = mx.from_fp8(weight, dtype=dtype)
+            # Per-tensor (rank 0) and per-expert ([E, 1, 1]) scales broadcast.
+            if broadcasts(scale_inv.shape, weight.shape):
+                return (weight * scale_inv).astype(dtype)
             bs = 128
             m, n = weight.shape
             pad_bottom = (-m) % bs
@@ -390,12 +399,14 @@ def dequant(weight, scale_inv):
         # Dequantize fp8
         new_weights = {}
         for k, v in weights.items():
-            if "weight_scale_inv" in k:
-                scale_inv = v
+            # Static activation scales have no consumer here.
+            if k.endswith("activation_scale"):
+                continue
+            # Expert scales are named "experts.down_proj_scale_inv", so match
+            # the suffix rather than a "weight_scale_inv" substring.
+            if k.endswith("_scale_inv"):
                 wk = k.replace("_scale_inv", "")
-                weight = weights[wk]
-                weight = dequant(weight, scale_inv)
-                new_weights[wk] = weight
+                new_weights[wk] = dequant(weights[wk], v)
             elif k not in new_weights:
                 new_weights[k] = v
         weights = new_weights

From 8931ba690b4dd285434f64aee9e9fac0d006c716 Mon Sep 17 00:00:00 2001
From: Michal Klein <46717574+michalk8@users.noreply.github.com>
Date: Thu, 17 Sep 2026 11:22:25 +0200
Subject: [PATCH 19/27] Fix Mistral4 test

---
 tests/test_models.py | 47 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 47 insertions(+)

diff --git a/tests/test_models.py b/tests/test_models.py
index 3740e1408..7ee980f86 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -2,6 +2,7 @@
 
 import copy
 import importlib
+import math
 import unittest
 from unittest import mock
 
@@ -1902,12 +1903,58 @@ def test_mistral4(self):
             max_position_embeddings=4096,
             rms_norm_eps=1e-6,
             first_k_dense_replace=0,
+            rope_parameters={
+                "rope_type": "yarn",
+                "factor": 128.0,
+                "mscale": 1.0,
+                "mscale_all_dim": 1.0,
+                "beta_fast": 32,
+                "beta_slow": 1,
+                "original_max_position_embeddings": 4096,
+                "rope_theta": 100000.0,
+                "llama_4_scaling_beta": 0.1,
+            },
         )
         model = mistral4_text.Model(args)
         self.model_test_runner(
             model, args.model_type, args.vocab_size, args.num_hidden_layers
         )
 
+        s = 0.1 * math.log(128.0) + 1.0
+        self.assertAlmostEqual(
+            model.layers[0].self_attn.scale,
+            (args.qk_nope_head_dim + args.qk_rope_head_dim) ** -0.5 * s * s,
+        )
+
+        # Real fp8 layout: rank-0 scales, [E, 1, 1] for expert stacks, no "weight_".
+        e, mi, h = args.n_routed_experts, args.moe_intermediate_size, args.hidden_size
+        fp8 = lambda *s: mx.full(s, 0x38, dtype=mx.uint8)  # 0x38 is 1.0 in e4m3
+        scale = lambda v, *s: mx.full(s, v, dtype=mx.bfloat16)
+        p = "model.layers.0.mlp"
+        out = model.sanitize(
+            {
+                f"{p}.shared_experts.down_proj.weight": fp8(h, mi),
+                f"{p}.shared_experts.down_proj.weight_scale_inv": scale(2.0),
+                f"{p}.shared_experts.down_proj.activation_scale": scale(1.0),
+                f"{p}.experts.down_proj": fp8(e, h, mi),
+                f"{p}.experts.down_proj_scale_inv": scale(4.0, e, 1, 1),
+                f"{p}.experts.down_proj_activation_scale": scale(1.0, e, 1, 1),
+                f"{p}.experts.gate_up_proj": fp8(e, 2 * mi, h),
+                f"{p}.experts.gate_up_proj_scale_inv": scale(8.0, e, 1, 1),
+            }
+        )
+        self.assertFalse([k for k in out if "activation_scale" in k])
+        self.assertFalse([k for k, v in out.items() if v.dtype == mx.uint8])
+        for name, shape, value in (
+            ("shared_experts.down_proj", (h, mi), 2.0),
+            ("switch_mlp.down_proj", (e, h, mi), 4.0),
+            ("switch_mlp.gate_proj", (e, mi, h), 8.0),
+            ("switch_mlp.up_proj", (e, mi, h), 8.0),
+        ):
+            w = out[f"{p}.{name}.weight"]
+            self.assertEqual(w.shape, shape)
+            self.assertEqual(w.reshape(-1)[0].item(), value)
+
     def test_gemma2(self):
         from mlx_lm.models import gemma2
 

From 7f684f0057574e400613e636d782bf6a6999c640 Mon Sep 17 00:00:00 2001
From: Michal Klein <46717574+michalk8@users.noreply.github.com>
Date: Thu, 17 Sep 2026 13:08:00 +0200
Subject: [PATCH 20/27] Use a linear layer for gate

---
 mlx_lm/models/mistral4_text.py | 35 +++++++++-------------------------
 1 file changed, 9 insertions(+), 26 deletions(-)

diff --git a/mlx_lm/models/mistral4_text.py b/mlx_lm/models/mistral4_text.py
index 7d39fb690..acf18108e 100644
--- a/mlx_lm/models/mistral4_text.py
+++ b/mlx_lm/models/mistral4_text.py
@@ -210,30 +210,6 @@ def mistral4_expert_select(
     return inds, selected_scores
 
 
-class Mistral4MoEGate(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.top_k = args.num_experts_per_tok
-        self.norm_topk_prob = args.norm_topk_prob
-        self.n_routed_experts = args.n_routed_experts
-        self.routed_scaling_factor = args.routed_scaling_factor
-        self.n_group = args.n_group
-        self.topk_group = args.topk_group
-        self.weight = mx.zeros((self.n_routed_experts, args.hidden_size))
-
-    def __call__(self, x):
-        gates = x @ self.weight.T
-        return mistral4_expert_select(
-            gates,
-            self.top_k,
-            self.n_group,
-            self.topk_group,
-            self.routed_scaling_factor,
-            self.norm_topk_prob,
-        )
-
-
 class Mistral4MoE(nn.Module):
     def __init__(self, args: ModelArgs):
         super().__init__()
@@ -245,7 +221,7 @@ def __init__(self, args: ModelArgs):
             args.n_routed_experts,
         )
 
-        self.gate = Mistral4MoEGate(args)
+        self.gate = nn.Linear(args.hidden_size, args.n_routed_experts, bias=False)
         if args.n_shared_experts is not None:
             intermediate_size = args.moe_intermediate_size * args.n_shared_experts
             self.shared_experts = DeepseekV3MLP(
@@ -253,7 +229,14 @@ def __init__(self, args: ModelArgs):
             )
 
     def __call__(self, x):
-        inds, scores = self.gate(x)
+        inds, scores = mistral4_expert_select(
+            self.gate(x),
+            self.num_experts_per_tok,
+            self.args.n_group,
+            self.args.topk_group,
+            self.args.routed_scaling_factor,
+            self.args.norm_topk_prob,
+        )
         y = self.switch_mlp(x, inds)
         y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
         if self.args.n_shared_experts is not None:

From b59538d5c496df3d7609e6bb94f10714c9584468 Mon Sep 17 00:00:00 2001
From: Michal Klein <46717574+michalk8@users.noreply.github.com>
Date: Thu, 17 Sep 2026 13:13:44 +0200
Subject: [PATCH 21/27] Fix Mistral4 pipelining

---
 mlx_lm/models/mistral4_text.py | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

diff --git a/mlx_lm/models/mistral4_text.py b/mlx_lm/models/mistral4_text.py
index acf18108e..59f9548e2 100644
--- a/mlx_lm/models/mistral4_text.py
+++ b/mlx_lm/models/mistral4_text.py
@@ -293,6 +293,9 @@ def __call__(
     ) -> mx.array:
         h = input_embeddings if input_embeddings is not None else self.embed_tokens(x)
 
+        pipeline_rank = self.pipeline_rank
+        pipeline_size = self.pipeline_size
+
         if cache is None:
             cache = [None] * len(self.pipeline_layers)
 
@@ -306,9 +309,23 @@ def __call__(
             self.args.rope_parameters["original_max_position_embeddings"],
         ).astype(h.dtype)
 
+        # Receive from the previous process in the pipeline
+        if pipeline_rank < pipeline_size - 1:
+            h = mx.distributed.recv_like(h, (pipeline_rank + 1))
+
         for l, c in zip(self.pipeline_layers, cache):
             h = l(h, attn_scale, mask, cache=c)
 
+        # Send to the next process in the pipeline
+        if pipeline_rank != 0:
+            h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
+            if cache[-1] is not None:
+                cache[-1].keys = mx.depends(cache[-1].keys, h)
+
+        # Broadcast h while keeping it in the graph
+        if pipeline_size > 1:
+            h = mx.distributed.all_gather(h)[: h.shape[0]]
+
         return self.norm(h)
 
 
@@ -336,7 +353,7 @@ def __call__(
 
     @property
     def layers(self):
-        return self.model.layers
+        return self.model.pipeline_layers
 
     def sanitize(self, weights):
         def broadcasts(scale_shape, weight_shape):

From ff4cb66814535ee0f05c832986586b10c88560ed Mon Sep 17 00:00:00 2001
From: Michal Klein <46717574+michalk8@users.noreply.github.com>
Date: Thu, 17 Sep 2026 13:25:03 +0200
Subject: [PATCH 22/27] Fix sharding

---
 mlx_lm/models/mistral4_text.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/mlx_lm/models/mistral4_text.py b/mlx_lm/models/mistral4_text.py
index 59f9548e2..972458765 100644
--- a/mlx_lm/models/mistral4_text.py
+++ b/mlx_lm/models/mistral4_text.py
@@ -5,7 +5,7 @@
 
 import mlx.core as mx
 import mlx.nn as nn
-from mlx.nn.layers.distributed import shard_inplace, shard_linear
+from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .deepseek_v3 import (
@@ -228,7 +228,12 @@ def __init__(self, args: ModelArgs):
                 args, intermediate_size=intermediate_size
             )
 
+        self.sharding_group = None
+
     def __call__(self, x):
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(x)
+
         inds, scores = mistral4_expert_select(
             self.gate(x),
             self.num_experts_per_tok,
@@ -241,6 +246,10 @@ def __call__(self, x):
         y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
         if self.args.n_shared_experts is not None:
             y = y + self.shared_experts(x)
+
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+
         return y
 
 
@@ -474,6 +483,8 @@ def shard(self, group: Optional[mx.distributed.Group] = None):
                 )
 
             else:
+                # Shard in place: the MoE aggregates the partial sums itself.
+                layer.mlp.sharding_group = group
                 if hasattr(layer.mlp, "shared_experts"):
                     shard_inplace(
                         layer.mlp.shared_experts.gate_proj,

From ba5b03cf9ced0b3a2829b459b868a7076c7040ff Mon Sep 17 00:00:00 2001
From: Michal Klein <46717574+michalk8@users.noreply.github.com>
Date: Thu, 17 Sep 2026 13:29:44 +0200
Subject: [PATCH 23/27] Don't materialize the mask

---
 mlx_lm/models/mistral4_text.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/mlx_lm/models/mistral4_text.py b/mlx_lm/models/mistral4_text.py
index 972458765..9eb702865 100644
--- a/mlx_lm/models/mistral4_text.py
+++ b/mlx_lm/models/mistral4_text.py
@@ -309,7 +309,10 @@ def __call__(
             cache = [None] * len(self.pipeline_layers)
 
         offset = cache[0].offset if cache[0] is not None else 0
-        mask = create_attention_mask(h, cache[0], return_array=True)
+        # No return_array: attention passes the mask straight to SDPA, so the
+        # fast "causal" path works. deepseek_v3 needs an array because its
+        # absorbed MLA path does mx.where(mask, pe_scores, ...).
+        mask = create_attention_mask(h, cache[0])
 
         attn_scale = _get_llama_4_attn_scale(
             x.shape[1],

From 27a8e01bb42a1b22cedde3b0489278b590d763f8 Mon Sep 17 00:00:00 2001
From: Michal Klein <46717574+michalk8@users.noreply.github.com>
Date: Thu, 17 Sep 2026 13:35:14 +0200
Subject: [PATCH 24/27] Rename mistral4_text to mistral4

---
 mlx_lm/models/mistral3.py                       | 6 +++---
 mlx_lm/models/{mistral4_text.py => mistral4.py} | 0
 mlx_lm/utils.py                                 | 1 -
 tests/test_models.py                            | 6 +++---
 4 files changed, 6 insertions(+), 7 deletions(-)
 rename mlx_lm/models/{mistral4_text.py => mistral4.py} (100%)

diff --git a/mlx_lm/models/mistral3.py b/mlx_lm/models/mistral3.py
index ba986990d..fa95e919c 100644
--- a/mlx_lm/models/mistral3.py
+++ b/mlx_lm/models/mistral3.py
@@ -6,7 +6,7 @@
 import mlx.core as mx
 import mlx.nn as nn
 
-from . import llama, ministral3, mistral4_text
+from . import llama, ministral3, mistral4
 from .base import BaseModelArgs
 
 
@@ -30,8 +30,8 @@ def __init__(self, args: ModelArgs):
                 ministral3.ModelArgs.from_dict(args.text_config)
             )
         elif args.text_config.get("model_type") == "mistral4":
-            self.language_model = mistral4_text.Model(
-                mistral4_text.ModelArgs.from_dict(args.text_config)
+            self.language_model = mistral4.Model(
+                mistral4.ModelArgs.from_dict(args.text_config)
             )
         else:
             self.language_model = llama.Model(
diff --git a/mlx_lm/models/mistral4_text.py b/mlx_lm/models/mistral4.py
similarity index 100%
rename from mlx_lm/models/mistral4_text.py
rename to mlx_lm/models/mistral4.py
diff --git a/mlx_lm/utils.py b/mlx_lm/utils.py
index f2fc10d89..4181dd826 100644
--- a/mlx_lm/utils.py
+++ b/mlx_lm/utils.py
@@ -45,7 +45,6 @@
 # Constants
 MODEL_REMAPPING = {
     "mistral": "llama",
-    "mistral4": "mistral4_text",
     "llava": "mistral3",
     "phi-msft": "phixtral",
     "falcon_mamba": "mamba",
diff --git a/tests/test_models.py b/tests/test_models.py
index 7ee980f86..8c2541666 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -1877,9 +1877,9 @@ def test_deepseek_v3(self):
         )
 
     def test_mistral4(self):
-        from mlx_lm.models import mistral4_text
+        from mlx_lm.models import mistral4
 
-        args = mistral4_text.ModelArgs(
+        args = mistral4.ModelArgs(
             model_type="mistral4",
             vocab_size=1024,
             hidden_size=128,
@@ -1915,7 +1915,7 @@ def test_mistral4(self):
                 "llama_4_scaling_beta": 0.1,
             },
         )
-        model = mistral4_text.Model(args)
+        model = mistral4.Model(args)
         self.model_test_runner(
             model, args.model_type, args.vocab_size, args.num_hidden_layers
         )

From 3ac79eefcc79ec6ec1c106c84a6b03d66bd7322c Mon Sep 17 00:00:00 2001
From: Michal Klein <46717574+michalk8@users.noreply.github.com>
Date: Thu, 17 Sep 2026 14:57:29 +0200
Subject: [PATCH 25/27] Use MultiLinear

---
 mlx_lm/models/mistral4.py | 122 ++++++++++++++++++++++++++++----------
 1 file changed, 90 insertions(+), 32 deletions(-)

diff --git a/mlx_lm/models/mistral4.py b/mlx_lm/models/mistral4.py
index 9eb702865..0a7cdddf6 100644
--- a/mlx_lm/models/mistral4.py
+++ b/mlx_lm/models/mistral4.py
@@ -13,6 +13,7 @@
     DeepseekV3Model,
 )
 from .ministral3 import _get_llama_4_attn_scale
+from .mla import MultiLinear
 from .pipeline import PipelineMixin
 from .rope_utils import apply_yarn_mscale, initialize_rope
 from .switch_layers import SwitchGLU
@@ -100,10 +101,12 @@ def __init__(self, args: ModelArgs):
             bias=args.attention_bias,
         )
         self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank, eps=1e-6)
-        self.kv_b_proj = nn.Linear(
-            self.kv_lora_rank,
-            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
-            bias=False,
+        # kv_b_proj absorbed, so the cache holds the compressed latent.
+        self.embed_q = MultiLinear(
+            self.qk_nope_head_dim, self.kv_lora_rank, self.num_heads
+        )
+        self.unembed_out = MultiLinear(
+            self.kv_lora_rank, self.v_head_dim, self.num_heads
         )
 
         self.o_proj = nn.Linear(
@@ -141,35 +144,52 @@ def __call__(
 
         compressed_kv = self.kv_a_proj_with_mqa(x)
         k_latent, k_rope = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
-
-        # Project latent to K and V
-        kv = self.kv_b_proj(self.kv_a_layernorm(k_latent))
-        kv = kv.reshape(B, L, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
-        kv = kv.transpose(0, 2, 1, 3)
-        k_nope, v = mx.split(kv, [self.qk_nope_head_dim], axis=-1)
-
         k_rope = k_rope.reshape(B, L, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3)
+        kv_latent = mx.expand_dims(self.kv_a_layernorm(k_latent), axis=1)
 
         offset = cache.offset if cache is not None else 0
         q_rope = self.rope(q_rope, offset)
         k_rope = self.rope(k_rope, offset)
 
-        # Expand k_rope to all heads
-        k_rope = mx.broadcast_to(
-            k_rope, [B, self.num_heads, k_rope.shape[2], self.qk_rope_head_dim]
-        )
-
-        query_states = mx.concatenate([q_nope, q_rope], axis=-1)
-        key_states = mx.concatenate([k_nope, k_rope], axis=-1)
-
-        query_states = query_states * attn_scale
+        # The llama-4 scale applies to the whole query, so scale both halves.
+        q_nope = q_nope * attn_scale
+        q_rope = q_rope * attn_scale
 
         if cache is not None:
-            key_states, v = cache.update_and_fetch(key_states, v)
-
-        output = scaled_dot_product_attention(
-            query_states, key_states, v, cache=cache, scale=self.scale, mask=mask
-        )
+            kv_latent, k_rope = cache.update_and_fetch(kv_latent, k_rope)
+
+        if L == 1:
+            # Decode: attend to the latent directly. pe_scores is [B, H, 1, L].
+            pe_scores = (q_rope * self.scale) @ k_rope.swapaxes(-1, -2)
+            if mask is not None:
+                pe_scores = mx.where(
+                    mask,
+                    pe_scores,
+                    mx.array(mx.finfo(pe_scores.dtype).min, pe_scores.dtype),
+                )
+            output = scaled_dot_product_attention(
+                self.embed_q(q_nope),
+                kv_latent,
+                kv_latent,
+                cache=cache,
+                scale=self.scale,
+                mask=pe_scores,
+            )
+            output = self.unembed_out(output)
+        else:
+            k_nope = self.embed_q(kv_latent, transpose=False)
+            v = self.unembed_out(kv_latent)
+            k_rope = mx.broadcast_to(
+                k_rope, [B, self.num_heads, k_rope.shape[2], self.qk_rope_head_dim]
+            )
+            output = scaled_dot_product_attention(
+                mx.concatenate([q_nope, q_rope], axis=-1),
+                mx.concatenate([k_nope, k_rope], axis=-1),
+                v,
+                cache=cache,
+                scale=self.scale,
+                mask=mask,
+            )
 
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
         return self.o_proj(output)
@@ -309,9 +329,7 @@ def __call__(
             cache = [None] * len(self.pipeline_layers)
 
         offset = cache[0].offset if cache[0] is not None else 0
-        # No return_array: attention passes the mask straight to SDPA, so the
-        # fast "causal" path works. deepseek_v3 needs an array because its
-        # absorbed MLA path does mx.where(mask, pe_scores, ...).
+        # "causal" suits prefill; at L == 1 it is None, which decode wants.
         mask = create_attention_mask(h, cache[0])
 
         attn_scale = _get_llama_4_attn_scale(
@@ -448,6 +466,41 @@ def dequant(weight, scale_inv):
                         ]
                         weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
 
+            # Absorb kv_b_proj into embed_q / unembed_out.
+            # TODO: affine only; non-affine modes have no biases key.
+            attn = f"{prefix}.self_attn"
+            if f"{attn}.kv_b_proj.weight" in weights:
+                quantized = f"{attn}.kv_b_proj.scales" in weights
+                w = weights.pop(f"{attn}.kv_b_proj.weight")
+                head_dim = self.args.qk_nope_head_dim + self.args.v_head_dim
+                if quantized:
+                    dims = self.args.kv_lora_rank
+                    scales = weights.pop(f"{attn}.kv_b_proj.scales")
+                    biases = weights.pop(f"{attn}.kv_b_proj.biases")
+                    bits = (w.shape[-1] * 32) // dims
+                    group_size = dims // scales.shape[-1]
+                    w = mx.dequantize(
+                        w, scales, biases, bits=bits, group_size=group_size
+                    )
+                w = w.reshape(self.args.num_attention_heads, head_dim, -1)
+                wk = mx.contiguous(
+                    w[:, : self.args.qk_nope_head_dim, :].swapaxes(-1, -2)
+                )
+                wv = mx.contiguous(w[:, self.args.qk_nope_head_dim :, :])
+                if quantized:
+                    wk, wk_scales, wk_biases = mx.quantize(
+                        wk, bits=bits, group_size=group_size
+                    )
+                    wv, wv_scales, wv_biases = mx.quantize(
+                        wv, bits=bits, group_size=group_size
+                    )
+                    weights[f"{attn}.embed_q.scales"] = wk_scales
+                    weights[f"{attn}.embed_q.biases"] = wk_biases
+                    weights[f"{attn}.unembed_out.scales"] = wv_scales
+                    weights[f"{attn}.unembed_out.biases"] = wv_biases
+                weights[f"{attn}.embed_q.weight"] = wk
+                weights[f"{attn}.unembed_out.weight"] = wv
+
         return {k: v for k, v in weights.items() if "rotary_emb.inv_freq" not in k}
 
     def shard(self, group: Optional[mx.distributed.Group] = None):
@@ -464,11 +517,16 @@ def shard(self, group: Optional[mx.distributed.Group] = None):
                     layer.self_attn.q_b_proj, "all-to-sharded", group=group
                 )
 
-            layer.self_attn.kv_b_proj = shard_linear(
-                layer.self_attn.kv_b_proj, "all-to-sharded", group=group
-            )
-
             layer.self_attn.num_heads //= N
+            num_heads = layer.self_attn.num_heads
+            sh = group.rank() * num_heads
+            eh = sh + num_heads
+
+            def shard_heads(w):
+                return w[sh:eh]
+
+            layer.self_attn.embed_q.apply(shard_heads)
+            layer.self_attn.unembed_out.apply(shard_heads)
 
             layer.self_attn.o_proj = shard_linear(
                 layer.self_attn.o_proj, "sharded-to-all", group=group

From 7bd8aa3633b8e6769aada30ec52f93379985e961 Mon Sep 17 00:00:00 2001
From: Michal Klein <46717574+michalk8@users.noreply.github.com>
Date: Thu, 17 Sep 2026 15:13:56 +0200
Subject: [PATCH 26/27] Simplify sanitize

---
 mlx_lm/models/mistral3.py | 16 ++++++----------
 1 file changed, 6 insertions(+), 10 deletions(-)

diff --git a/mlx_lm/models/mistral3.py b/mlx_lm/models/mistral3.py
index fa95e919c..aa4167f34 100644
--- a/mlx_lm/models/mistral3.py
+++ b/mlx_lm/models/mistral3.py
@@ -49,19 +49,15 @@ def __call__(
         )
 
     def sanitize(self, weights):
-        sanitized = {}
+        lm_weights = {}
         for key, value in weights.items():
             if "vision_tower" in key or "multi_modal_projector" in key:
                 continue
-            if key.startswith("model."):
-                key = key[len("model.") :]
-            sanitized[key] = value
-
-        lm_weights = {
-            k[len("language_model.") :]: v
-            for k, v in sanitized.items()
-            if k.startswith("language_model.")
-        }
+            if key.startswith("model.language_model."):
+                key = "model." + key.removeprefix("model.language_model.")
+            else:
+                key = key.removeprefix("language_model.")
+            lm_weights[key] = value
 
         sanitized_lm = self.language_model.sanitize(lm_weights)
         return {"language_model." + k: v for k, v in sanitized_lm.items()}

From c6b40355f1d1d888c1bc826c8a6c3197f3591b5c Mon Sep 17 00:00:00 2001
From: Michal Klein <46717574+michalk8@users.noreply.github.com>
Date: Thu, 17 Sep 2026 15:36:24 +0200
Subject: [PATCH 27/27] Use h.shape[1]

---
 mlx_lm/models/mistral4.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mlx_lm/models/mistral4.py b/mlx_lm/models/mistral4.py
index 0a7cdddf6..572060c31 100644
--- a/mlx_lm/models/mistral4.py
+++ b/mlx_lm/models/mistral4.py
@@ -333,7 +333,7 @@ def __call__(
         mask = create_attention_mask(h, cache[0])
 
         attn_scale = _get_llama_4_attn_scale(
-            x.shape[1],
+            h.shape[1],
             offset,
             self.args.rope_parameters["llama_4_scaling_beta"],
             self.args.rope_parameters["original_max_position_embeddings"],