Fix Gemma3n inference without cache (#323)

Closes #322

mlx_lm/LEARNED_QUANTS.md

@@ -22,7 +22,7 @@ quantized model can be further refined with DWQ.
 To get started, first install the requirements:
 
 ```
-pip install mlx-lm[quant]
+pip install "mlx-lm[train]"
 ```
 
 ### DWQ

mlx_lm/LORA.md

@@ -26,6 +26,12 @@ LoRA (QLoRA).[^qlora] LoRA fine-tuning works with the following model families:
 
 ## Run
 
+First, make sure you have the training dependenices installed:
+
+```shell
+pip install "mlx-lm[train]"
+```
+
 The main command is `mlx_lm.lora`. To see a full list of command-line options run:
 
 ```shell

mlx_lm/_version.py

@@ -1,3 +1,3 @@
-# Copyright © 2023-2024 Apple Inc.
+# Copyright © 2023-2025 Apple Inc.
 
-__version__ = "0.26.0"
+__version__ = "0.26.2"

mlx_lm/chat.py

@@ -69,6 +69,11 @@ def setup_arg_parser():
         default=DEFAULT_MAX_TOKENS,
         help="Maximum number of tokens to generate",
     )
+    parser.add_argument(
+        "--system-prompt",
+        default=None,
+        help="System prompt to be used for the chat template",
+    )
     return parser
 
 
@@ -104,7 +109,10 @@ def main():
         if query == "h":
             print_help()
             continue
-        messages = [{"role": "user", "content": query}]
+        messages = []
+        if args.system_prompt is not None:
+            messages.append({"role": "system", "content": args.system_prompt})
+        messages.append({"role": "user", "content": query})
         prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
         for response in stream_generate(
             model,

mlx_lm/evaluate.py

@@ -20,6 +20,7 @@ import mlx.nn as nn
 import numpy as np
 from lm_eval.api.model import LM
 from lm_eval.api.registry import register_model
+from lm_eval.models import huggingface
 from tqdm import tqdm
 
 from .generate import stream_generate
@@ -62,7 +63,7 @@ def chat_template_fn(**extra_kwargs):
 @register_model("mlxlm")
 class MLXLM(LM):
 
-    tokenizer_name = lm_eval.models.huggingface.HFLM.tokenizer_name
+    tokenizer_name = huggingface.HFLM.tokenizer_name
     apply_chat_template = chat_template_fn()
 
     def __init__(

mlx_lm/examples/afm/README.md

@@ -1,10 +0,0 @@
-Apple Foundation Model in MLX
-=============================
-
-This example provides information about porting the AFM model to MLX-LM and
-training adapters with it or using it as any other open-weights model. It is
-paired with https://developer.apple.com/apple-intelligence/foundation-models-adapter/ that
-was published during WWDC 25 and to get the weights one needs to follow these
-instructions to download the toolkit.
-
-

mlx_lm/examples/afm/extract_quantized.py

@@ -1,38 +0,0 @@
-import argparse
-
-import mlx.core as mx
-
-from mlx_lm.convert import convert
-
-
-def mixed_quant(layer_path, layer, cfg):
-    if "embedding" in layer_path:
-        return {"group_size": 32, "bits": 8}
-    return hasattr(layer, "to_quantized")
-
-
-def main(argv):
-    parser = argparse.ArgumentParser(
-        description="Quantize the AFM according to its original quantization"
-    )
-    parser.add_argument("source", help="The mlx model containing the fp32 weights")
-    parser.add_argument("destination", help="The folder to save the quantized model to")
-    parser.add_argument("--copy-adapters", action="store_true")
-    parser.add_argument(
-        "--dtype", choices=["bfloat16", "float16", "float32"], default="float32"
-    )
-    args = parser.parse_args(argv)
-
-    convert(
-        args.source,
-        args.destination,
-        quantize=True,
-        q_group_size=128,
-        q_bits=2,
-        dtype=getattr(mx, args.dtype),
-        quant_predicate=mixed_quant,
-    )
-
-
-if __name__ == "__main__":
-    main(None)

mlx_lm/examples/afm/extract_weights.py

@@ -1,249 +0,0 @@
-import argparse
-import json
-import textwrap
-from pathlib import Path
-
-import torch
-from safetensors.torch import save_file
-from transformers import LlamaTokenizerFast
-
-
-def share_data(a, b):
-    return a.untyped_storage().data_ptr() == b.untyped_storage().data_ptr()
-
-
-def get_model_config():
-    return {
-        "model_type": "afm7",
-        "vocab_size": 153600,
-        "hidden_dim": 2048,
-        "num_layers": 56,
-        "num_kv_reuse_layers": 21,
-        "num_heads": 16,
-        "num_kv_heads": 2,
-        "hidden_dim_scale_factor": 3.25,
-        "rope_theta": 500000.0,
-    }
-
-
-def get_adapter_config():
-    return {
-        "num_layers": 56,
-        "lora_parameters": {
-            "rank": 32,
-            "scale": 0.5,
-            "dropout": 0.0,
-            "keys": [
-                "mlp.gate_proj",
-                "mlp.down_proj",
-                "mlp.up_proj",
-                "self_attn.qkv_proj",
-                "self_attn.q_proj",
-                "self_attn.out_proj",
-            ],
-        },
-    }
-
-
-def get_chat_template():
-    return textwrap.dedent(
-        """
-        {%- set default_system_message = "A conversation between a user and a helpful assistant." %}
-        {%- if messages[0]['role'] == 'system' %}
-            {%- set system_message = messages[0]['content'] %}
-            {%- set loop_messages = messages[1:] %}
-        {%- else %}
-            {%- set system_message = default_system_message %}
-            {%- set loop_messages = messages %}
-        {%- endif %}
-        {{- '<turn_start> system<n>' + system_message -}}
-        {% if tools %}
-            {{- ('<n>system tools: ' + (tools | map('tojson') | join('<n>'))) -}}
-        {% endif %}
-        {{- '<turn_end>' -}}
-        {% for message in loop_messages %}
-            {{- '<turn_start> ' + message['role'] + '<n>' + message['content'] + '<turn_end>' -}}
-        {% endfor %}
-        {% if add_generation_prompt is defined and add_generation_prompt %}
-            {% if messages[-1]['role'] != 'assistant' %}
-                {{- '<turn_start> assistant<n>' -}}
-            {% endif %}
-        {% endif %}"""
-    ).strip()
-
-
-def map_model_keys(state):
-    model_keys = {}
-    for old in state:
-        if "adapter" in old:
-            continue
-        if "kv_quantizer" in old:
-            continue
-
-        new = old
-        if new.startswith("layers."):
-            new = new[7:]
-            new = new.replace("layer_", "")
-            new = new.replace("attention.norm", "input_layernorm")
-            new = new.replace(".attention.", ".self_attn.")
-            new = new.replace("self_attn.output_transform", "self_attn.out_proj")
-            new = new.replace("feed_forward.norm", "post_attention_layernorm")
-            new = new.replace(".feed_forward.", ".mlp.")
-            new = new.replace("hidden_transform.linear_0", "gate_proj")
-            new = new.replace("hidden_transform.linear_1", "up_proj")
-            new = new.replace("mlp.output_transform", "mlp.down_proj")
-            if new.startswith("segment_0"):
-                new = new.replace("segment_0", "layers")
-                new = new.replace(".qkv_transform.", ".qkv_proj.")
-                new = new.replace(".fused_linear.", ".")
-                new = new.replace(".qk_norm.query_norm.", ".q_norm.")
-                new = new.replace(".qk_norm.key_norm.", ".k_norm.")
-            elif new.startswith("segment_1"):
-                new = new.replace("segment_1", "kv_reuse_layers")
-                new = new.replace(".q_transform.", ".q_proj.")
-                new = new.replace(".q_norm.query_norm.", ".q_norm.")
-        new = new.replace(".wrapped.", ".")
-        new = "model." + new
-        model_keys[old] = new
-    return model_keys
-
-
-def map_adapter_keys(state):
-    adapter_keys = {}
-    for old in state:
-        if "adapter" not in old:
-            continue
-
-        new = old
-        new = new[7:]
-        new = new.replace("layer_", "")
-        new = new.replace(".attention.", ".self_attn.")
-        new = new.replace("self_attn.output_transform", "self_attn.out_proj")
-        new = new.replace(".feed_forward.", ".mlp.")
-        new = new.replace("hidden_transform.linear_0", "gate_proj")
-        new = new.replace("hidden_transform.linear_1", "up_proj")
-        new = new.replace("mlp.output_transform", "mlp.down_proj")
-        if new.startswith("segment_0"):
-            new = new.replace("segment_0", "layers")
-            new = new.replace(".qkv_transform.", ".qkv_proj.")
-            new = new.replace(".fused_linear.", ".")
-        elif new.startswith("segment_1"):
-            new = new.replace("segment_1", "kv_reuse_layers")
-            new = new.replace(".q_transform.", ".q_proj.")
-
-        new = new.replace(".lora_0.b_transpose", ".b_transpose.0")
-        new = new.replace(".lora_1.b_transpose", ".b_transpose.1")
-        new = new.replace(".lora_2.b_transpose", ".b_transpose.2")
-        new = new.replace(".lora_0.a_transpose", ".a_transpose.0")
-        new = new.replace(".lora_1.a_transpose", ".a_transpose.1")
-        new = new.replace(".lora_2.a_transpose", ".a_transpose.2")
-        new = new.replace("adapters.base_adapter.b_transpose", "lora_b")
-        new = new.replace("adapters.base_adapter.a_transpose", "lora_a")
-        new = "model." + new
-        adapter_keys[old] = new
-    return adapter_keys
-
-
-def add_kv_quant_weights(new_state, old_state, dt):
-    for k, v in old_state.items():
-        if "range" not in k:
-            continue
-
-        v = v.tolist()
-        weight = "quant_key_scale" if "key_quantizer" in k else "quant_value_scale"
-        new_k = k[: k.find("kv_quantizer")]
-        new_k = new_k.replace("segment_0.layer_", "")
-        new_k = new_k.replace("attention", "self_attn")
-        new_k = "model." + new_k + weight
-        quant_scale = torch.tensor(max(v[0] / (-128), v[1] / 127), dtype=dt)
-        new_state[new_k] = quant_scale
-
-
-def cast(x, dt):
-    info = torch.finfo(dt)
-    a, b = info.min, info.max
-    return x.clip(a, b).to(dt)
-
-
-def main(argv):
-    parser = argparse.ArgumentParser(
-        description="Map the PT weights to MLX-LM safetensors"
-    )
-    parser.add_argument("source", help="The source weights in PT format")
-    parser.add_argument("tokenizer", help="The source tokenizer file")
-    parser.add_argument("destination", help="The folder to write the model weights in")
-    parser.add_argument(
-        "--dtype", choices=["bfloat16", "float16", "float32"], default="float32"
-    )
-    parser.add_argument(
-        "--adapter-dtype", choices=["bfloat16", "float16", "float32"], default="float32"
-    )
-    parser.add_argument(
-        "--force",
-        "-f",
-        action="store_true",
-        help="If set overwrite the weight files in the destination folder",
-    )
-    args = parser.parse_args(argv)
-
-    destination = Path(args.destination)
-    if not destination.exists():
-        destination.mkdir()
-    model_file = destination / "model.safetensors"
-    adapter_file = destination / "adapters.safetensors"
-    if (model_file.exists() or adapter_file.exists()) and not args.force:
-        print("Model files already exist. Delete them or use --force to overwrite them")
-        return
-
-    # Write the configuration files
-    with (destination / "config.json").open("w") as f:
-        json.dump(get_model_config(), f, indent=4)
-    with (destination / "adapter_config.json").open("w") as f:
-        json.dump(get_adapter_config(), f, indent=4)
-
-    # Pop the tied output transform
-    state = torch.load(args.source)
-    if share_data(state["embedding.weight"], state["output_transform.weight"]):
-        state.pop("output_transform.weight")
-
-    # Map the weights
-    model_keys = map_model_keys(state)
-    adapter_keys = map_adapter_keys(state)
-
-    # Make the new weight dictionaries
-    dt = getattr(torch, args.dtype)
-    adapter_dt = getattr(torch, args.adapter_dtype)
-    adapters = {
-        k_new: cast(state[k_old], adapter_dt) for k_old, k_new in adapter_keys.items()
-    }
-    model = {k_new: cast(state[k_old], dt) for k_old, k_new in model_keys.items()}
-    add_kv_quant_weights(model, state, dt)
-
-    # Save them to disk
-    save_file(model, model_file)
-    save_file(adapters, adapter_file)
-
-    # Save the tokenizer
-    tok = LlamaTokenizerFast(vocab_file=args.tokenizer)
-    tok.chat_template = get_chat_template()
-    tok.eos_token_ids = tok.convert_tokens_to_ids("<turn_end>")
-    tok.save_pretrained(str(destination))
-    with (destination / "tokenizer_config.json").open("r+") as f:
-        config = json.load(f)
-        config["tokenizer_class"] = "NewlineTokenizer"
-        f.seek(0)
-        json.dump(config, f, indent=4)
-        f.truncate()
-    with (destination / "tokenizer.json").open("r+") as f:
-        tok = json.load(f)
-        tok["decoder"]["decoders"].insert(
-            1,
-            {"type": "Replace", "pattern": {"String": "<n>"}, "content": "\n"},
-        )
-        f.seek(0)
-        json.dump(tok, f, indent=4)
-        f.truncate()
-
-
-if __name__ == "__main__":
-    main(None)

mlx_lm/examples/afm/requirements.txt

@@ -1,3 +0,0 @@
-tamm==0.1.0
-transformers
-torch

mlx_lm/examples/pipeline_generate.py

@@ -50,7 +50,7 @@ def shard_and_load(repo):
 
     # Lazy load and shard model to figure out
     # which weights we need
-    model, _ = load_model(model_path, lazy=True, strict=False)
+    model, config = load_model(model_path, lazy=True, strict=False)
 
     group = mx.distributed.init()
     rank = group.rank()
@@ -68,7 +68,11 @@ def shard_and_load(repo):
     download(args.model, allow_patterns=local_files)
 
     # Load and shard the model, and load the weights
-    tokenizer = load_tokenizer(model_path)
+    tokenizer = load_tokenizer(
+        model_path,
+        {"trust_remote_code": True},
+        eos_token_ids=config.get("eos_token_id", None),
+    )
     model, _ = load_model(model_path, lazy=True, strict=False)
     model.model.pipeline(group)
     mx.eval(model.parameters())

mlx_lm/generate.py

@@ -218,31 +218,31 @@ def wired_limit(model: nn.Module, streams: Optional[List[mx.Stream]] = None):
         try:
             yield
         finally:
-            return
-
-    model_bytes = tree_reduce(
-        lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, model, 0
-    )
-    max_rec_size = mx.metal.device_info()["max_recommended_working_set_size"]
-    if model_bytes > 0.9 * max_rec_size:
-        model_mb = model_bytes // 2**20
-        max_rec_mb = max_rec_size // 2**20
-        print(
-            f"[WARNING] Generating with a model that requires {model_mb} MB "
-            f"which is close to the maximum recommended size of {max_rec_mb} "
-            "MB. This can be slow. See the documentation for possible work-arounds: "
-            "https://github.com/ml-explore/mlx-lm/tree/main#large-models"
+            pass
+    else:
+        model_bytes = tree_reduce(
+            lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, model, 0
         )
-    old_limit = mx.set_wired_limit(max_rec_size)
-    try:
-        yield
-    finally:
-        if streams is not None:
-            for s in streams:
-                mx.synchronize(s)
-        else:
-            mx.synchronize()
-        mx.set_wired_limit(old_limit)
+        max_rec_size = mx.metal.device_info()["max_recommended_working_set_size"]
+        if model_bytes > 0.9 * max_rec_size:
+            model_mb = model_bytes // 2**20
+            max_rec_mb = max_rec_size // 2**20
+            print(
+                f"[WARNING] Generating with a model that requires {model_mb} MB "
+                f"which is close to the maximum recommended size of {max_rec_mb} "
+                "MB. This can be slow. See the documentation for possible work-arounds: "
+                "https://github.com/ml-explore/mlx-lm/tree/main#large-models"
+            )
+        old_limit = mx.set_wired_limit(max_rec_size)
+        try:
+            yield
+        finally:
+            if streams is not None:
+                for s in streams:
+                    mx.synchronize(s)
+            else:
+                mx.synchronize()
+            mx.set_wired_limit(old_limit)
 
 
 @dataclass
@@ -329,21 +329,25 @@ def generate_step(
            when ``kv_bits`` is non-None. Default: ``0``.
         prompt_progress_callback (Callable[int, int]): A call-back which takes the
            prompt tokens processed so far and the total number of prompt tokens.
-        input_embeddings (mx.array, optional): Input embeddings to use in conjunction
-           with prompt tokens. Default: ``None``.
+        input_embeddings (mx.array, optional): Input embeddings to use instead of or in
+          conjunction with prompt tokens. Default: ``None``.
 
     Yields:
         Tuple[mx.array, mx.array]: One token and a vector of log probabilities.
     """
-    if len(prompt) == 0:
-        raise ValueError("Prompt must be non-empty.")
     if input_embeddings is not None:
         if not does_model_support_input_embeddings(model):
             raise ValueError("Model does not support input embeddings.")
-        elif prompt.shape[0] != input_embeddings.shape[0]:
+        elif len(prompt) > 0 and len(prompt) != len(input_embeddings):
             raise ValueError(
-                "If using input embeddings, the sequence length must match that of the prompt."
+                f"When providing input_embeddings, their sequence length ({len(input_embeddings)}) "
+                f"must match the sequence length of the prompt ({len(prompt)}), or the "
+                "prompt must be empty."
             )
+    elif len(prompt) == 0:
+        raise ValueError(
+            "Either input_embeddings or prompt (or both) must be provided."
+        )
 
     tokens = None
 
@@ -386,7 +390,7 @@ def generate_step(
 
             logits = logits[:, -1, :]
 
-            if logits_processors:
+            if logits_processors and len(input_tokens) > 0:
                 tokens = (
                     mx.concat([tokens, input_tokens])
                     if tokens is not None
@@ -402,7 +406,9 @@ def generate_step(
             return sampled, logprobs.squeeze(0)
 
     with mx.stream(generation_stream):
-        total_prompt_tokens = prompt.shape[0]
+        total_prompt_tokens = (
+            len(input_embeddings) if input_embeddings is not None else len(prompt)
+        )
         prompt_processed_tokens = 0
         while total_prompt_tokens - prompt_processed_tokens > prefill_step_size:
             _model_call(

mlx_lm/lora.py

@@ -46,6 +46,9 @@ CONFIG_DEFAULTS = {
     "optimizer_config": {
         "adam": {},
         "adamw": {},
+        "muon": {},
+        "sgd": {},
+        "adafactor": {},
     },
     "data": "data/",
     "seed": 0,
@@ -103,9 +106,9 @@ def build_parser():
     parser.add_argument(
         "--optimizer",
         type=str,
-        choices=["adam", "adamw"],
+        choices=["adam", "adamw", "sgd", "adafactor"],
         default=None,
-        help="Optimizer to use for training: adam or adamw",
+        help="Optimizer to use for training: adam, adamw, sgd, or adafactor.",
     )
     parser.add_argument(
         "--mask-prompt",
@@ -251,11 +254,16 @@ def train_model(
 
     optimizer_name = args.optimizer.lower()
     optimizer_config = args.optimizer_config.get(optimizer_name, {})
-
     if optimizer_name == "adam":
         opt_class = optim.Adam
     elif optimizer_name == "adamw":
         opt_class = optim.AdamW
+    elif optimizer_name == "muon":
+        opt_class = optim.Muon
+    elif optimizer_name == "sgd":
+        opt_class = optim.SGD
+    elif optimizer_name == "adafactor":
+        opt_class = optim.Adafactor
     else:
         raise ValueError(f"Unsupported optimizer: {optimizer_name}")
 

mlx_lm/models/cache.py

@@ -453,9 +453,9 @@ class RotatingKVCache(_BaseCache):
         raise NotImplementedError("RotatingKVCache Quantization NYI")
 
 
-class MambaCache(_BaseCache):
-    def __init__(self):
-        self.cache = [None, None]
+class ArraysCache(_BaseCache):
+    def __init__(self, size):
+        self.cache = [None] * size
 
     def __setitem__(self, idx, value):
         self.cache[idx] = value
@@ -472,6 +472,11 @@ class MambaCache(_BaseCache):
         self.cache = v
 
 
+class MambaCache(ArraysCache):
+    def __init__(self):
+        super().__init__(size=2)
+
+
 class ChunkedKVCache(KVCache):
     def __init__(self, chunk_size=None):
         super().__init__()

mlx_lm/models/deepseek_v3.py

@@ -303,7 +303,9 @@ def group_expert_select(
     group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True)
     k = n_group - topk_group
     group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :]
-    scores = mx.put_along_axis(scores, group_idx, mx.array(0.0), axis=-2)
+    scores = mx.put_along_axis(
+        scores, mx.stop_gradient(group_idx), mx.array(0.0), axis=-2
+    )
     scores = mx.flatten(scores, -2, -1)
 
     k = top_k

mlx_lm/models/exaone4.py

@@ -0,0 +1,208 @@
+# Copyright © 2023-2024 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import KVCache, RotatingKVCache
+from .rope_utils import initialize_rope
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    rms_norm_eps: float
+    vocab_size: int
+    num_key_value_heads: int
+    max_position_embeddings: int
+    rope_theta: float
+    head_dim: int
+    tie_word_embeddings: bool
+    rope_scaling: Dict[str, Union[float, str]]
+    sliding_window: Optional[int]
+    sliding_window_pattern: Optional[str]
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs, is_local: Optional[bool]):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        assert args.num_key_value_heads is not None
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+
+        head_dim = args.head_dim
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
+
+        self.q_norm = nn.RMSNorm(head_dim, eps=args.rms_norm_eps)
+        self.k_norm = nn.RMSNorm(head_dim, eps=args.rms_norm_eps)
+        self.is_local = is_local or False
+        self.use_rope = is_local is None or is_local
+        if self.use_rope:
+            self.rope = initialize_rope(
+                head_dim,
+                base=args.rope_theta,
+                traditional=False,
+                scaling_config=args.rope_scaling,
+                max_position_embeddings=args.max_position_embeddings,
+            )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        queries = self.q_norm(queries.reshape(B, L, self.n_heads, -1)).transpose(
+            0, 2, 1, 3
+        )
+        keys = self.k_norm(keys.reshape(B, L, self.n_kv_heads, -1)).transpose(
+            0, 2, 1, 3
+        )
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            if self.use_rope:
+                queries = self.rope(queries, offset=cache.offset)
+                keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        elif self.use_rope:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs, is_local: bool):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.hidden_size = args.hidden_size
+        self.self_attn = Attention(args, is_local)
+        self.mlp = MLP(args.hidden_size, args.intermediate_size)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+        self.post_feedforward_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+        self.args = args
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(x, mask, cache)
+        h = x + self.post_attention_layernorm(r)
+        r = self.mlp(h)
+        out = h + self.post_feedforward_layernorm(r)
+        return out
+
+
+class ExaoneModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.num_hidden_layers = args.num_hidden_layers
+        assert self.vocab_size > 0
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        pattern = args.sliding_window_pattern
+        self.layers = [
+            TransformerBlock(
+                args=args,
+                is_local=pattern[i % len(pattern)] == "L" if pattern else None,
+            )
+            for i in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        h = self.embed_tokens(inputs)
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = ExaoneModel(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,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        out = self.model(inputs, mask, cache)
+        if self.args.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
+        return out
+
+    def make_cache(self):
+        return [
+            (
+                RotatingKVCache(max_size=self.args.sliding_window, keep=0)
+                if l.self_attn.is_local
+                else KVCache()
+            )
+            for l in self.layers
+        ]
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/gemma3n.py

@@ -25,7 +25,6 @@ class TextConfig(BaseModelArgs):
     vocab_size: int
     num_key_value_heads: int
     num_kv_shared_layers: int
-    query_pre_attn_scalar: float
     vocab_size_per_layer_input: int
     sliding_window: int
     max_position_embeddings: int
@@ -177,7 +176,11 @@ class MLP(nn.Module):
         super().__init__()
         self.config = config
         self.hidden_size = config.hidden_size
-        self.intermediate_size = config.intermediate_size
+        self.intermediate_size = (
+            config.intermediate_size[layer_idx]
+            if isinstance(config.intermediate_size, list)
+            else config.intermediate_size
+        )
         self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
         self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
         self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
@@ -473,7 +476,7 @@ class LanguageModel(nn.Module):
         per_layer_inputs = self.project_per_layer_inputs(h, per_layer_inputs)
 
         if cache is None:
-            cache = [None] * len(self.layers)
+            cache = self.make_cache()
 
         if mask is None:
             full_mask = create_attention_mask(

mlx_lm/models/glm4_moe.py

@@ -0,0 +1,329 @@
+# Copyright © 2025 Apple Inc.
+
+import math
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, Dict, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    max_position_embeddings: int
+    moe_intermediate_size: int
+    norm_topk_prob: bool
+    num_attention_heads: int
+    n_group: int
+    head_dim: int
+    topk_group: int
+    n_shared_experts: int
+    n_routed_experts: int
+    routed_scaling_factor: float
+    num_experts_per_tok: int
+    first_k_dense_replace: int
+    num_hidden_layers: int
+    num_key_value_heads: int
+    rms_norm_eps: float
+    rope_theta: float
+    rope_scaling: Optional[Dict]
+    use_qk_norm: bool
+    tie_word_embeddings: bool
+    attention_bias: bool
+    partial_rotary_factor: float
+    scoring_func: str = "sigmoid"
+    topk_method: str = "noaux_tc"
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+
+        head_dim = args.head_dim
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=args.attention_bias)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.attention_bias)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.attention_bias)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
+
+        self.use_qk_norm = args.use_qk_norm
+        if self.use_qk_norm:
+            self.q_norm = nn.RMSNorm(head_dim, eps=args.rms_norm_eps)
+            self.k_norm = nn.RMSNorm(head_dim, eps=args.rms_norm_eps)
+
+        self.rope = nn.RoPE(
+            int(head_dim * args.partial_rotary_factor),
+            traditional=False,
+            base=args.rope_theta,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        queries = queries.reshape(B, L, self.n_heads, -1)
+        keys = keys.reshape(B, L, self.n_kv_heads, -1)
+
+        if self.use_qk_norm:
+            queries = self.q_norm(queries)
+            keys = self.k_norm(keys)
+
+        queries = queries.transpose(0, 2, 1, 3)
+        keys = keys.transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(
+        self, config: ModelArgs, hidden_size: int = None, intermediate_size: int = None
+    ):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size if hidden_size is None else hidden_size
+        self.intermediate_size = (
+            config.intermediate_size if intermediate_size is None else intermediate_size
+        )
+
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+
+    def __call__(self, x):
+        down_proj = self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+        return down_proj
+
+
+@mx.compile
+def group_expert_select(
+    gates,
+    e_score_correction_bias,
+    top_k,
+    n_group,
+    topk_group,
+    routed_scaling_factor,
+    norm_topk_prob,
+):
+
+    scores = mx.sigmoid(gates.astype(mx.float32))
+    orig_scores = scores
+    scores = scores + e_score_correction_bias
+    if n_group > 1:
+        k = top_k
+        scores = mx.unflatten(scores, axis=-1, shape=(n_group, -1))
+        group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True)
+        k = n_group - topk_group
+        group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :]
+        scores = mx.put_along_axis(
+            scores, mx.stop_gradient(group_idx), mx.array(0.0), axis=-2
+        )
+        scores = mx.flatten(scores, -2, -1)
+
+    k = top_k
+    inds = mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k]
+    scores = mx.take_along_axis(orig_scores, inds, axis=-1)
+    if top_k > 1 and norm_topk_prob:
+        denominator = scores.sum(axis=-1, keepdims=True)
+        scores = scores / denominator
+    scores = scores * routed_scaling_factor
+
+    return inds, scores
+
+
+class MoEGate(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.top_k = config.num_experts_per_tok
+        self.norm_topk_prob = config.norm_topk_prob
+        self.n_routed_experts = config.n_routed_experts
+        self.routed_scaling_factor = config.routed_scaling_factor
+        self.n_group = config.n_group
+        self.topk_group = config.topk_group
+        self.weight = mx.zeros((self.n_routed_experts, config.hidden_size))
+        self.e_score_correction_bias = mx.zeros((self.n_routed_experts,))
+        assert config.topk_method == "noaux_tc", "Unsupported topk method."
+
+    def __call__(self, x):
+        return group_expert_select(
+            x @ self.weight.T,
+            self.e_score_correction_bias,
+            self.top_k,
+            self.n_group,
+            self.topk_group,
+            self.routed_scaling_factor,
+            self.norm_topk_prob,
+        )
+
+
+class MoE(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.num_experts_per_tok = config.num_experts_per_tok
+        self.switch_mlp = SwitchGLU(
+            config.hidden_size,
+            config.moe_intermediate_size,
+            config.n_routed_experts,
+        )
+
+        self.gate = MoEGate(config)
+        if config.n_shared_experts is not None:
+            intermediate_size = config.moe_intermediate_size * config.n_shared_experts
+            self.shared_experts = MLP(
+                config=config, 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.config.n_shared_experts is not None:
+            y = y + self.shared_experts(x)
+
+        return y
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, config: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.self_attn = Attention(config)
+        self.mlp = (
+            MoE(config)
+            if (
+                config.n_routed_experts is not None
+                and layer_idx >= config.first_k_dense_replace
+            )
+            else MLP(config)
+        )
+        self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            config.hidden_size, eps=config.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 LanguageModel(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = [
+            DecoderLayer(config, idx) for idx in range(config.num_hidden_layers)
+        ]
+        self.start_idx = 0
+        self.end_idx = len(self.layers)
+        self.num_layers = self.end_idx
+
+        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def __call__(
+        self,
+        x: mx.array,
+        cache: Optional[Any] = None,
+        mask: Optional[mx.array] = None,
+    ) -> mx.array:
+        h = self.embed_tokens(x)
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * self.num_layers
+
+        for i in range(self.num_layers):
+            h = self.layers[self.start_idx + i](h, mask, cache[i])
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.args = config
+        self.model_type = config.model_type
+        self.model = LanguageModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+        mask: Optional[mx.array] = None,
+    ):
+        out = self.model(inputs, cache, mask)
+        return self.lm_head(out)
+
+    def sanitize(self, weights):
+        mpt_layer = self.args.num_hidden_layers
+
+        # Stack experts
+        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)
+
+        # Remove multi-token prediction layer
+        return {
+            k: v
+            for k, v in weights.items()
+            if not k.startswith(f"model.layers.{mpt_layer}")
+        }
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    @property
+    def cast_predicate(self):
+        def predicate(k):
+            return "e_score_correction_bias" not in k
+
+        return predicate

mlx_lm/models/lfm2.py

@@ -0,0 +1,278 @@
+# Copyright © 2025 Apple Inc.
+from dataclasses import dataclass
+from typing import Any, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import ArraysCache, KVCache
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    max_position_embeddings: int
+    norm_eps: float
+    conv_bias: bool
+    conv_L_cache: int
+    block_dim: int
+    block_ff_dim: int
+    block_multiple_of: int
+    block_ffn_dim_multiplier: float
+    block_auto_adjust_ff_dim: bool
+    full_attn_idxs: List[int]
+    rope_theta: float
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+
+        self.head_dim = head_dim = args.hidden_size // n_heads
+
+        self.scale = head_dim**-0.5
+
+        self.q_layernorm = nn.RMSNorm(head_dim, eps=args.norm_eps)
+        self.k_layernorm = nn.RMSNorm(head_dim, eps=args.norm_eps)
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
+        self.out_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
+
+        self.rope = nn.RoPE(
+            self.head_dim,
+            base=args.rope_theta,
+            traditional=False,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        queries = self.q_layernorm(queries.reshape(B, L, self.n_heads, -1)).transpose(
+            0, 2, 1, 3
+        )
+        keys = self.k_layernorm(keys.reshape(B, L, self.n_kv_heads, -1)).transpose(
+            0, 2, 1, 3
+        )
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, mask=mask, scale=self.scale
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.out_proj(output)
+
+
+class ShortConv(nn.Module):
+    def __init__(
+        self,
+        args: ModelArgs,
+        layer_idx: int,
+    ):
+        super().__init__()
+        self.args = args
+        self.layer_idx = layer_idx
+        self.L_cache = args.conv_L_cache
+        self.bias = args.conv_bias
+
+        self.conv = nn.Conv1d(
+            in_channels=args.hidden_size,
+            out_channels=args.hidden_size,
+            kernel_size=self.L_cache,
+            groups=args.hidden_size,
+            bias=self.bias,
+        )
+        self.in_proj = nn.Linear(args.hidden_size, 3 * args.hidden_size, bias=self.bias)
+        self.out_proj = nn.Linear(args.hidden_size, args.hidden_size, bias=self.bias)
+
+    def __call__(
+        self,
+        x: mx.array,
+        cache: Optional[Any] = None,
+    ):
+        seqlen = x.shape[1]
+        BCx = self.in_proj(x)
+        B, C, x = mx.split(BCx, 3, axis=-1)
+        Bx = B * x
+
+        state = None
+        if cache is not None:
+            state = cache[0]
+        if state is None:
+            state = mx.zeros(
+                (Bx.shape[0], self.L_cache - 1, self.args.hidden_size), dtype=Bx.dtype
+            )
+
+        Bx = mx.concatenate([state, Bx], axis=-2)
+        if cache is not None:
+            cache[0] = Bx[:, -(self.L_cache - 1) :]
+        conv_out = self.conv(Bx)
+
+        y = C * conv_out
+        return self.out_proj(y)
+
+
+class MLP(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        ff_dim: int,
+        multiple_of: int,
+        auto_adjust_ff_dim: bool,
+        ffn_dim_multiplier: Optional[float],
+    ):
+        super().__init__()
+        if auto_adjust_ff_dim:
+            ff_dim = int(2 * ff_dim / 3)
+            if ffn_dim_multiplier is not None:
+                ff_dim = int(ffn_dim_multiplier * ff_dim)
+            ff_dim = multiple_of * ((ff_dim + multiple_of - 1) // multiple_of)
+
+        self.w1 = nn.Linear(dim, ff_dim, bias=False)
+        self.w3 = nn.Linear(dim, ff_dim, bias=False)
+        self.w2 = nn.Linear(ff_dim, dim, bias=False)
+
+    def __call__(self, x) -> mx.array:
+        return self.w2(nn.silu(self.w1(x)) * self.w3(x))
+
+
+class Lfm2DecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.is_attention_layer = layer_idx in args.full_attn_idxs
+
+        if self.is_attention_layer:
+            self.self_attn = Attention(args)
+        else:
+            self.conv = ShortConv(args, layer_idx)
+        self.feed_forward = MLP(
+            dim=args.block_dim,
+            ff_dim=args.block_ff_dim,
+            multiple_of=args.block_multiple_of,
+            auto_adjust_ff_dim=args.block_auto_adjust_ff_dim,
+            ffn_dim_multiplier=args.block_ffn_dim_multiplier,
+        )
+
+        self.operator_norm = nn.RMSNorm(args.hidden_size, eps=args.norm_eps)
+        self.ffn_norm = nn.RMSNorm(args.hidden_size, eps=args.norm_eps)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+
+        if self.is_attention_layer:
+            r = self.self_attn(self.operator_norm(x), mask=mask, cache=cache)
+        else:
+            r = self.conv(
+                self.operator_norm(x),
+                cache=cache,
+            )
+        h = x + r
+        out = h + self.feed_forward(self.ffn_norm(h))
+        return out
+
+
+class Lfm2Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.num_hidden_layers = args.num_hidden_layers
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            Lfm2DecoderLayer(args, layer_idx=i) for i in range(args.num_hidden_layers)
+        ]
+
+        self.embedding_norm = nn.RMSNorm(args.hidden_size, eps=args.norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        if input_embeddings is not None:
+            h = input_embeddings
+        else:
+            h = self.embed_tokens(inputs)
+
+        if mask is None:
+            first_attn_idx = self.args.full_attn_idxs[0]
+            c = [cache[first_attn_idx]] if cache is not None else None
+            mask = create_attention_mask(h, c)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, cache=c)
+
+        return self.embedding_norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = Lfm2Model(args)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        out = self.model(inputs, mask, cache, input_embeddings)
+        return self.model.embed_tokens.as_linear(out)
+
+    def sanitize(self, weights):
+        sanitized_weights = {}
+        for name, param in weights.items():
+            if "conv.weight" in name:
+                if param.shape[-1] > param.shape[1]:
+                    param = param.transpose(0, 2, 1)
+
+            sanitized_weights[name] = param
+        return sanitized_weights
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def make_cache(self):
+        return [
+            KVCache() if l.is_attention_layer else ArraysCache(size=1)
+            for l in self.layers
+        ]

mlx_lm/quant/dwq.py

@@ -30,8 +30,9 @@ class Catcher(nn.Module):
         self.module = module
 
     def __call__(self, *args, **kwargs):
-        self.outputs = self.module(*args, **kwargs)
-        return self.outputs
+        outputs = self.module(*args, **kwargs)
+        self.outputs = outputs[0] if isinstance(outputs, tuple) else outputs
+        return outputs
 
 
 def dwq_quantize(
@@ -54,6 +55,7 @@ def dwq_quantize(
         if hasattr(m, "bits") and hasattr(m, "group_size"):
             m.unfreeze(keys=["scales", "biases"], recurse=False)
 
+    q_model.train()
     q_model.apply_to_modules(unfreeze)
     print_trainable_parameters(q_model)
 
@@ -213,15 +215,19 @@ def main():
     mx.random.seed(args.seed)
 
     model_path, hf_repo = get_model_path(args.model, revision=None)
-    model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
+    model, config, tokenizer = fetch_from_hub(
+        model_path, lazy=True, trust_remote_code=True
+    )
 
     calibration_data = load_data(
         tokenizer, args.data_path, args.num_samples, args.max_seq_length
     )
 
     if args.quantized_model is not None:
-        q_model_path = get_model_path(args.quantized_model, revision=None)
-        q_model, config, _ = fetch_from_hub(q_model_path, lazy=True)
+        q_model_path, _ = get_model_path(args.quantized_model, revision=None)
+        q_model, config, _ = fetch_from_hub(
+            q_model_path, lazy=True, trust_remote_code=True
+        )
     else:
         q_model = copy.deepcopy(model)
         _, config = quantize_model(

mlx_lm/sample_utils.py

@@ -16,7 +16,7 @@ def make_sampler(
     xtc_probability: float = 0.0,
     xtc_threshold: float = 0.0,
     xtc_special_tokens: List[int] = [],
-) -> Callable[mx.array, mx.array]:
+) -> Callable[[mx.array], mx.array]:
     """
     Make a sampler function for use with ``generate_step``.
 

mlx_lm/server.py

@@ -737,6 +737,9 @@ class APIHandler(BaseHTTPRequestHandler):
 
         self.prompt_cache.tokens.extend(tokens)
 
+        if gen_response.finish_reason is not None:
+            finish_reason = gen_response.finish_reason
+
         logging.debug(f"Prompt: {gen_response.prompt_tps:.3f} tokens-per-sec")
         logging.debug(f"Generation: {gen_response.generation_tps:.3f} tokens-per-sec")
         logging.debug(f"Peak memory: {gen_response.peak_memory:.3f} GB")
@@ -839,7 +842,7 @@ class APIHandler(BaseHTTPRequestHandler):
         """
         Respond to a GET request from a client.
         """
-        if self.path == "/v1/models":
+        if self.path.startswith("/v1/models"):
             self.handle_models_request()
         elif self.path == "/health":
             self.handle_health_check()
@@ -867,11 +870,18 @@ class APIHandler(BaseHTTPRequestHandler):
 
         files = ["config.json", "model.safetensors.index.json", "tokenizer_config.json"]
 
+        parts = self.path.split("/")
+        filter_repo_id = None
+        if len(parts) > 3:
+            filter_repo_id = "/".join(parts[3:])
+
         def probably_mlx_lm(repo):
             if repo.repo_type != "model":
                 return False
             if "main" not in repo.refs:
                 return False
+            if filter_repo_id is not None and repo.repo_id != filter_repo_id:
+                return False
             file_names = {f.file_path.name for f in repo.refs["main"].files}
             return all(f in file_names for f in files)
 

mlx_lm/tokenizer_utils.py

@@ -89,7 +89,7 @@ class NaiveStreamingDetokenizer(StreamingDetokenizer):
     def text(self):
         if self._current_tokens:
             self._current_text = self._tokenizer.decode(self._current_tokens)
-            if (
+            if self._current_text.endswith("\ufffd") or (
                 self._tokenizer.clean_up_tokenization_spaces
                 and len(self._current_text) > 0
                 and self._current_text[-1] == " "

mlx_lm/tuner/callbacks.py

@@ -32,12 +32,19 @@ class WandBCallback(TrainingCallback):
         self.wrapped_callback = wrapped_callback
         wandb.init(project=project_name, dir=log_dir, config=config)
 
+    def _convert_to_serializable(self, data: dict) -> dict:
+        return {k: v.tolist() if hasattr(v, "tolist") else v for k, v in data.items()}
+
     def on_train_loss_report(self, train_info: dict):
-        wandb.log(train_info, step=train_info.get("iteration"))
+        wandb.log(
+            self._convert_to_serializable(train_info), step=train_info.get("iteration")
+        )
         if self.wrapped_callback:
             self.wrapped_callback.on_train_loss_report(train_info)
 
     def on_val_loss_report(self, val_info: dict):
-        wandb.log(val_info, step=val_info.get("iteration"))
+        wandb.log(
+            self._convert_to_serializable(val_info), step=val_info.get("iteration")
+        )
         if self.wrapped_callback:
             self.wrapped_callback.on_val_loss_report(val_info)

mlx_lm/tuner/trainer.py

@@ -102,13 +102,14 @@ def iterate_batches(
 
     # If running in distributed mode (N machines) then each one should skip N-1
     # samples
+    offset = mx.distributed.init().rank()
     step = mx.distributed.init().size()
     if batch_size % step != 0:
         raise ValueError("The batch size must be divisible by the number of workers")
 
     # Make the batches:
     batch_idx = [
-        idx[i : i + batch_size : step]
+        idx[i + offset : i + offset + batch_size : step]
         for i in range(0, len(idx) - batch_size + 1, batch_size)
     ]
 
@@ -196,7 +197,8 @@ def train(
     iterate_batches: callable = iterate_batches,
     training_callback: TrainingCallback = None,
 ):
-    mx.set_wired_limit(mx.metal.device_info()["max_recommended_working_set_size"])
+    if mx.metal.is_available():
+        mx.set_wired_limit(mx.metal.device_info()["max_recommended_working_set_size"])
     print(f"Starting training..., iters: {args.iters}")
     world = mx.distributed.init()
     world_size = world.size()

mlx_lm/tuner/utils.py

@@ -88,6 +88,7 @@ def linear_to_lora_layers(
         "mistral",
         "mistral3",
         "llama",
+        "lfm2",
         "phi",
         "mixtral",
         "nemotron",
@@ -119,6 +120,7 @@ def linear_to_lora_layers(
         "ernie4_5",
         "dots1",
         "smollm3",
+        "exaone4",
     }:
         keys = {"self_attn.q_proj", "self_attn.v_proj"}
         if model.model_type in ["mixtral", "phimoe"]:

mlx_lm/utils.py

@@ -41,9 +41,10 @@ from .tuner.utils import get_total_parameters, load_adapters
 
 # Constants
 MODEL_REMAPPING = {
-    "mistral": "llama",  # mistral is compatible with llama
+    "mistral": "llama",
     "phi-msft": "phixtral",
     "falcon_mamba": "mamba",
+    "kimi_k2": "deepseek_v3",
 }
 
 MAX_FILE_SIZE_GB = 5
@@ -79,7 +80,9 @@ def compute_bits_per_weight(model):
     return model_bytes * 8 / model_params
 
 
-def get_model_path(path_or_hf_repo: str, revision: Optional[str] = None) -> Path:
+def get_model_path(
+    path_or_hf_repo: str, revision: Optional[str] = None
+) -> Tuple[Path, Optional[str]]:
     """
     Ensures the model is available locally. If the path does not exist locally,
     it is downloaded from the Hugging Face Hub.
@@ -140,7 +143,7 @@ def load_model(
     strict: bool = True,
     model_config: dict = {},
     get_model_classes: Callable[[dict], Tuple[Type[nn.Module], Type]] = _get_classes,
-) -> nn.Module:
+) -> Tuple[nn.Module, dict]:
     """
     Load and initialize the model from a given path.
 
@@ -158,7 +161,7 @@ def load_model(
             Defaults to the ``_get_classes`` function.
 
     Returns:
-        nn.Module: The loaded and initialized model.
+        Tuple[nn.Module, dict[str, Any]]: The loaded and initialized model and config.
 
     Raises:
         FileNotFoundError: If the weight files (.safetensors) are not found.
@@ -450,7 +453,7 @@ def quantize_model(
     quant_predicate: Optional[
         Callable[[str, nn.Module, dict], Union[bool, dict]]
     ] = None,
-) -> Tuple:
+) -> Tuple[nn.Module, dict]:
     """
     Applies quantization to the model weights.
 

requirements.txt

@@ -1,4 +1,4 @@
-mlx>=0.25.0
+mlx>=0.26.0
 numpy
 transformers>=4.39.3
 protobuf

setup.py

@@ -15,7 +15,7 @@ from _version import __version__
 setup(
     name="mlx-lm",
     version=__version__,
-    description="LLMs on Apple silicon with MLX and the Hugging Face Hub",
+    description="LLMs with MLX and the Hugging Face Hub",
     long_description=open("README.md", encoding="utf-8").read(),
     long_description_content_type="text/markdown",
     readme="README.md",
@@ -28,8 +28,8 @@ setup(
     python_requires=">=3.8",
     extras_require={
         "test": ["datasets"],
+        "train": ["datasets", "tqdm"],
         "evaluate": ["lm-eval", "tqdm"],
-        "quant": ["datasets", "tqdm"],
     },
     entry_points={
         "console_scripts": [

tests/test_chat.py

@@ -0,0 +1,167 @@
+import argparse
+import unittest
+from unittest.mock import MagicMock, patch
+
+from mlx_lm.chat import setup_arg_parser
+
+
+class TestChat(unittest.TestCase):
+
+    def test_setup_arg_parser_system_prompt(self):
+        parser = setup_arg_parser()
+
+        # Test default (no system prompt)
+        args = parser.parse_args([])
+        self.assertIsNone(args.system_prompt)
+
+        # Test with system prompt
+        args = parser.parse_args(["--system-prompt", "You are a helpful assistant."])
+        self.assertEqual(args.system_prompt, "You are a helpful assistant.")
+
+    def test_setup_arg_parser_all_args(self):
+        parser = setup_arg_parser()
+        args = parser.parse_args(
+            [
+                "--model",
+                "test-model",
+                "--adapter-path",
+                "/path/to/adapter",
+                "--temp",
+                "0.7",
+                "--top-p",
+                "0.9",
+                "--xtc-probability",
+                "0.1",
+                "--xtc-threshold",
+                "0.1",
+                "--seed",
+                "42",
+                "--max-kv-size",
+                "1024",
+                "--max-tokens",
+                "512",
+                "--system-prompt",
+                "You are a helpful assistant.",
+            ]
+        )
+
+        self.assertEqual(args.model, "test-model")
+        self.assertEqual(args.adapter_path, "/path/to/adapter")
+        self.assertEqual(args.temp, 0.7)
+        self.assertEqual(args.top_p, 0.9)
+        self.assertEqual(args.xtc_probability, 0.1)
+        self.assertEqual(args.xtc_threshold, 0.1)
+        self.assertEqual(args.seed, 42)
+        self.assertEqual(args.max_kv_size, 1024)
+        self.assertEqual(args.max_tokens, 512)
+        self.assertEqual(args.system_prompt, "You are a helpful assistant.")
+
+    @patch("mlx_lm.chat.load")
+    @patch("mlx_lm.chat.make_prompt_cache")
+    @patch("mlx_lm.chat.stream_generate")
+    @patch("builtins.input")
+    @patch("builtins.print")
+    def test_system_prompt_in_messages(
+        self,
+        mock_print,
+        mock_input,
+        mock_stream_generate,
+        mock_make_prompt_cache,
+        mock_load,
+    ):
+        from mlx_lm.chat import main
+
+        # Mock the model and tokenizer
+        mock_model = MagicMock()
+        mock_tokenizer = MagicMock()
+        mock_tokenizer.apply_chat_template.return_value = "processed_prompt"
+        mock_load.return_value = (mock_model, mock_tokenizer)
+
+        # Mock prompt cache
+        mock_prompt_cache = MagicMock()
+        mock_make_prompt_cache.return_value = mock_prompt_cache
+
+        # Mock stream_generate to return some responses
+        mock_response = MagicMock()
+        mock_response.text = "Hello there!"
+        mock_stream_generate.return_value = [mock_response]
+
+        # Mock user input: first a question, then 'q' to quit
+        mock_input.side_effect = ["What is the weather?", "q"]
+
+        # Test with system prompt
+        with patch(
+            "sys.argv", ["chat.py", "--system-prompt", "You are a weather assistant."]
+        ):
+            try:
+                main()
+            except SystemExit:
+                pass
+
+        # Verify that apply_chat_template was called with system prompt
+        mock_tokenizer.apply_chat_template.assert_called()
+        call_args = mock_tokenizer.apply_chat_template.call_args[0][
+            0
+        ]  # First positional arg (messages)
+
+        # Check that the messages contain both system and user messages
+        self.assertEqual(len(call_args), 2)
+        self.assertEqual(call_args[0]["role"], "system")
+        self.assertEqual(call_args[0]["content"], "You are a weather assistant.")
+        self.assertEqual(call_args[1]["role"], "user")
+        self.assertEqual(call_args[1]["content"], "What is the weather?")
+
+    @patch("mlx_lm.chat.load")
+    @patch("mlx_lm.chat.make_prompt_cache")
+    @patch("mlx_lm.chat.stream_generate")
+    @patch("builtins.input")
+    @patch("builtins.print")
+    def test_no_system_prompt_in_messages(
+        self,
+        mock_print,
+        mock_input,
+        mock_stream_generate,
+        mock_make_prompt_cache,
+        mock_load,
+    ):
+        from mlx_lm.chat import main
+
+        # Mock the model and tokenizer
+        mock_model = MagicMock()
+        mock_tokenizer = MagicMock()
+        mock_tokenizer.apply_chat_template.return_value = "processed_prompt"
+        mock_load.return_value = (mock_model, mock_tokenizer)
+
+        # Mock prompt cache
+        mock_prompt_cache = MagicMock()
+        mock_make_prompt_cache.return_value = mock_prompt_cache
+
+        # Mock stream_generate to return some responses
+        mock_response = MagicMock()
+        mock_response.text = "Hello there!"
+        mock_stream_generate.return_value = [mock_response]
+
+        # Mock user input: first a question, then 'q' to quit
+        mock_input.side_effect = ["What is the weather?", "q"]
+
+        # Test without system prompt
+        with patch("sys.argv", ["chat.py"]):
+            try:
+                main()
+            except SystemExit:
+                pass
+
+        # Verify that apply_chat_template was called without system prompt
+        mock_tokenizer.apply_chat_template.assert_called()
+        call_args = mock_tokenizer.apply_chat_template.call_args[0][
+            0
+        ]  # First positional arg (messages)
+
+        # Check that the messages contain only user message
+        self.assertEqual(len(call_args), 1)
+        self.assertEqual(call_args[0]["role"], "user")
+        self.assertEqual(call_args[0]["content"], "What is the weather?")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_models.py

@@ -251,6 +251,33 @@ class TestModels(unittest.TestCase):
             model, args.model_type, args.vocab_size, args.num_hidden_layers
         )
 
+    def test_lfm2(self):
+        from mlx_lm.models import lfm2
+
+        args = lfm2.ModelArgs(
+            model_type="lfm2",
+            hidden_size=1024,
+            num_hidden_layers=4,
+            num_attention_heads=4,
+            num_key_value_heads=2,
+            norm_eps=1e-5,
+            vocab_size=10_000,
+            full_attn_idxs=[0, 1, 2],
+            rope_theta=10000,
+            block_dim=1024,
+            block_ffn_dim_multiplier=1.5,
+            block_auto_adjust_ff_dim=True,
+            block_ff_dim=2048,
+            block_multiple_of=256,
+            max_position_embeddings=1000,
+            conv_bias=True,
+            conv_L_cache=3,
+        )
+        model = lfm2.Model(args)
+        self.model_test_runner(
+            model, args.model_type, args.vocab_size, args.num_hidden_layers
+        )
+
     def test_bitnet(self):
         from mlx_lm.models import bitnet
 

tests/test_tokenizers.py

@@ -48,6 +48,9 @@ class TestTokenizers(unittest.TestCase):
         tokens = tokenizer.encode("こんにちは！私の名前はAI")
         check(tokens)
 
+        tokens = tokenizer.encode("⊕ ⊻ ∧ ¬")
+        check(tokens)
+
         tokens = tokenizer.encode("a ,b")
         check(tokens)
 

tests/test_tuner_trainer.py

@@ -0,0 +1,66 @@
+# Copyright © 2025 Apple Inc.
+
+import unittest
+
+import mlx.core as mx
+
+from mlx_lm.tuner.trainer import iterate_batches
+
+
+class MockDistributedGroup:
+    def __init__(self, rank, size):
+        self._rank = rank
+        self._size = size
+
+    def rank(self):
+        return self._rank
+
+    def size(self):
+        return self._size
+
+
+class MockDistributed:
+    def __init__(self):
+        self.rank = 0
+        self.size = 1
+
+    def init(self):
+        return MockDistributedGroup(self.rank, self.size)
+
+
+class TestTunerTrainer(unittest.TestCase):
+    def test_iterate_batches_ddp(self):
+        olddist = mx.distributed
+        try:
+            mx.distributed = MockDistributed()
+
+            def run(rank, size, batch):
+                mx.distributed.rank = rank
+                mx.distributed.size = size
+
+                data = mx.arange(128).reshape(-1, 1).tolist()
+                data = [(d, 0) for d in data]
+
+                samples = set()
+                for i, (b, l) in enumerate(iterate_batches(data, batch, 1)):
+                    samples.add(tuple(mx.flatten(b).tolist()))
+
+                ref_batches = mx.arange(128).reshape(-1, batch).tolist()
+                for b in ref_batches:
+                    self.assertTrue(tuple(b[rank::size]) in samples)
+
+            run(0, 1, 4)
+            run(0, 1, 8)
+            run(0, 2, 8)
+            run(1, 2, 8)
+            run(0, 4, 8)
+            run(1, 4, 8)
+            run(2, 4, 8)
+            run(3, 4, 8)
+
+        finally:
+            mx.distributed = olddist
+
+
+if __name__ == "__main__":
+    unittest.main()