bases model can't be quantized

* Allow loading already quantized model + nits

* updates in eval

* add template args

* speed up eval

* fix

* shuffle data based on seed

ACKNOWLEDGMENTS.md

@@ -9,4 +9,4 @@ MLX LM was developed with contributions from the following individuals:
 
 - Shunta Saito: Added support for PLaMo models.
 - Prince Canuma: Helped add support for `Starcoder2` models.
-- Gökdeniz Gülmez: Added support for the following architectures: OpenBMB's `MiniCPM` and `MiniCPM3`, Kyutai's `Helium`, State-Space's`Mamba v1`, and Allenai's `OLMoE`; Added support for the following training algorithms: `full-fine-tuning`.
+- Gökdeniz Gülmez: Added support for the following architectures: OpenBMB's `MiniCPM` and `MiniCPM3`, Kyutai's `Helium`, State-Space's`Mamba v1`, Z.ai & THUKEG's `GLM4`, and Allenai's `OLMoE`; Added support for the following training algorithms: `full-fine-tuning`; Added support for the following other features: `Multiple Optimizers to choose for training`.

mlx_lm/LEARNED_QUANTS.md

@@ -0,0 +1,93 @@
+# Learned Quantization 
+
+To reduce the quality loss from quantization MLX LM has two options:
+
+- Distilled Weight Quantization (DWQ)
+- Activation-aware Weight Quantization (AWQ)[^1].
+
+Both DWQ and AWQ use an example dataset to tune parameters of the model. DWQ
+fine-tunes non-quantized parameters (including quantization scales and biases)
+using the non-quantized model as a teacher. AWQ scales and clips the weights
+prior to quantization. The scaling and clipping values are found with a grid
+search minimizing the distance from the quantized hidden activations to the
+non-quantized hidden activations
+
+To get started, first install the requirements:
+
+```
+pip install mlx-lm[lwq]
+```
+
+### DWQ
+
+Use `mlx_lm.dwq` to run DWQ on a given model. For example:
+
+```bash
+mlx_lm.dwq --model mistralai/Mistral-7B-Instruct-v0.3
+```
+
+Some important options, along with their default values are:
+
+- `--mlx-path mlx_model`: The location to save the DWQ model.
+- `--bits 4`: Precision of the quantization.
+- `--num-samples 1024`: Number of samples to use. Using more samples can lead to
+  better results but takes longer.
+- `--batch-size 8`: Use a smaller batch size to reduce the memory footprint.
+
+For a full list of options run:
+
+```bash
+mlx_lm.dwq --help
+```
+
+### AWQ 
+
+Use `mlx_lm.awq` to run AWQ on a given model. For example:
+
+```bash
+mlx_lm.awq --model mistralai/Mistral-7B-Instruct-v0.3
+```
+
+The script can take anywhere form a few minutes to several hours to run
+depending on the model size and the number of samples.
+
+Some important options, along with their default values, are:
+
+- `--mlx-path mlx_model`: The location to save the AWQ model.
+- `--bits 4`: Precision of the quantization.
+- `--num-samples 32`: Number of samples to use. Using more samples can lead to
+  better results but takes longer.
+- `--n-grid 10`: The granularity of the AWQ search. A larger grid can lead to
+  better results but takes longer.
+
+For a full list of options run:
+
+```bash
+mlx_lm.awq --help
+```
+
+### Evaluate
+
+Once the training script finishes, you can evaluate the quality of the model
+on downstream tasks using `mlx_lm.evaluate`. For example:
+
+```bash
+mlx_lm.evaluate \
+    --model mlx_model \
+    --tasks winogrande boolq arc_challenge arc_easy hellaswag openbookqa piqa social_iqa                     
+```
+
+### Upload to Hugging Face
+
+Use `mlx_lm.upload` to upload the quantized model to the Hugging Face Hub. For
+example:
+
+```bash
+mlx_lm.upload \
+    --path mlx_model \
+    --upload-repo mlx-community/Mistral-7B-Instruct-v0.3-3bit-DWQ
+```
+
+[^1]: Refer to the [paper](https://arxiv.org/abs/2306.00978)
+and [github repository](https://github.com/mit-han-lab/llm-awq) for more
+details.

mlx_lm/LORA.md

@@ -291,7 +291,7 @@ example:
 
 ```yaml
 hf_dataset:
-  name: "billsum"
+  path: "billsum"
   prompt_feature: "text"
   completion_feature: "summary"
 ```
@@ -308,12 +308,12 @@ with the same structure as above. For example:
 
 ```yaml
 hf_dataset:
-  - name: "Open-Orca/OpenOrca"
+  - path: "Open-Orca/OpenOrca"
     train_split: "train[:90%]"
     valid_split: "train[-10%:]"
     prompt_feature: "question"
     completion_feature: "response"
-  - name: "trl-lib/ultrafeedback_binarized"
+  - path: "trl-lib/ultrafeedback_binarized"
     train_split: "train[:90%]"
     valid_split: "train[-10%:]"
     chat_feature: "chosen"

mlx_lm/SERVER.md

@@ -86,6 +86,13 @@ curl localhost:8080/v1/chat/completions \
 - `adapters`: (Optional) A string path to low-rank adapters. The path must be
   relative to the directory the server was started in.
 
+- `draft_model`: (Optional) Specifies a smaller model to use for speculative
+  decoding. Set to `null` to unload.
+
+- `num_draft_tokens`: (Optional) The number of draft tokens the draft model
+  should predict at once. Defaults to `3`.
+
+
 ### Response Fields
 
 - `id`: A unique identifier for the chat.

mlx_lm/__main__.py

@@ -4,7 +4,21 @@ import importlib
 import sys
 
 if __name__ == "__main__":
-    subcommands = {"convert"}
+    subcommands = {
+        "awq",
+        "dwq",
+        "cache_prompt",
+        "chat",
+        "convert",
+        "evaluate",
+        "fuse",
+        "generate",
+        "lora",
+        "merge",
+        "server",
+        "manage",
+        "upload",
+    }
     if len(sys.argv) < 2:
         raise ValueError(f"CLI requires a subcommand in {subcommands}")
     subcommand = sys.argv.pop(1)

mlx_lm/_version.py

@@ -1,3 +1,3 @@
 # Copyright © 2023-2024 Apple Inc.
 
-__version__ = "0.22.2"
+__version__ = "0.24.0"

mlx_lm/awq.py

@@ -0,0 +1,605 @@
+# Copyright © 2025 Apple Inc.
+
+import argparse
+import copy
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Callable, Dict
+from urllib import request
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten, tree_map, tree_map_with_path
+from tqdm import tqdm
+
+from mlx_lm.models.base import create_attention_mask
+from mlx_lm.models.switch_layers import SwitchLinear
+from mlx_lm.utils import (
+    fetch_from_hub,
+    get_model_path,
+    save,
+)
+
+
+@dataclass
+class ScaleConfig:
+    prev: nn.Module
+    layers: list[nn.Module]
+    block: nn.Module | None = None
+    kwargs: list = field(default_factory=list)
+    use_config: Callable[[nn.Module], bool] | None = None
+
+
+@dataclass
+class AWQConfig:
+    embed: str
+    lm_head: str
+    no_clip: list[str]
+    scale_configs: list[ScaleConfig]
+    lm_key: str | None = None
+
+
+def update(cfg, **kwargs):
+    cfg = copy.deepcopy(cfg)
+    for k, v in kwargs.items():
+        setattr(cfg, k, v)
+    return cfg
+
+
+llama_awq = AWQConfig(
+    embed="embed_tokens",
+    lm_head="lm_head",
+    no_clip=["q_proj", "k_proj"],
+    scale_configs=[
+        ScaleConfig(
+            block="self_attn",
+            prev="input_layernorm",
+            layers=["q_proj", "k_proj", "v_proj"],
+            kwargs=["mask"],
+        ),
+        ScaleConfig(prev="mlp.up_proj", layers=["mlp.down_proj"]),
+        ScaleConfig(
+            block="mlp",
+            prev="post_attention_layernorm",
+            layers=["gate_proj", "up_proj"],
+        ),
+    ],
+)
+
+gemma3_text_awq = AWQConfig(
+    embed="embed_tokens",
+    lm_head="lm_head",
+    no_clip=["q_proj", "k_proj"],
+    scale_configs=[
+        ScaleConfig(
+            block="self_attn",
+            prev="input_layernorm",
+            layers=["q_proj", "k_proj", "v_proj"],
+            kwargs=["mask"],
+        ),
+        ScaleConfig(prev="mlp.up_proj", layers=["mlp.down_proj"]),
+        ScaleConfig(
+            block="mlp",
+            prev="pre_feedforward_layernorm",
+            layers=["gate_proj", "up_proj"],
+        ),
+    ],
+)
+
+gemma3_awq = update(gemma3_text_awq, lm_key="language_model")
+
+deepseek_v2_awq = AWQConfig(
+    embed="embed_tokens",
+    lm_head="lm_head",
+    no_clip=["q_proj", "q_a_proj", "q_b_proj", "kv_a_proj_with_mqa", "kv_b_proj"],
+    scale_configs=[
+        ScaleConfig(
+            block="self_attn",
+            prev="input_layernorm",
+            layers=["q_proj", "kv_a_proj_with_mqa"],
+            kwargs=["mask"],
+        ),
+        ScaleConfig(
+            prev="self_attn.kv_a_layernorm",
+            layers=["self_attn.kv_b_proj"],
+        ),
+        ScaleConfig(
+            prev="mlp.up_proj",
+            layers=["mlp.down_proj"],
+            use_config=lambda block: not "switch_mlp" in block.mlp,
+        ),
+        ScaleConfig(
+            prev="mlp.shared_experts.up_proj",
+            layers=["mlp.shared_experts.down_proj"],
+            use_config=lambda block: "switch_mlp" in block.mlp,
+        ),
+        ScaleConfig(
+            prev="mlp.switch_mlp.up_proj",
+            layers=["mlp.switch_mlp.down_proj"],
+            use_config=lambda block: "switch_mlp" in block.mlp,
+            kwargs=["indices"],
+        ),
+        ScaleConfig(
+            block="mlp",
+            prev="post_attention_layernorm",
+            layers=["gate_proj", "up_proj"],
+            use_config=lambda block: not "switch_mlp" in block.mlp,
+        ),
+        ScaleConfig(
+            block="mlp",
+            prev="post_attention_layernorm",
+            layers=[
+                "switch_mlp.gate_proj",
+                "switch_mlp.up_proj",
+                "shared_experts.gate_proj",
+                "shared_experts.up_proj",
+                "gate",  # not quantized, just scaled
+            ],
+            use_config=lambda block: "switch_mlp" in block.mlp,
+        ),
+    ],
+)
+
+AWQ_MODEL_CONFIGS = {
+    "llama": llama_awq,
+    "mistral": llama_awq,
+    "qwen2": llama_awq,
+    "qwen3": llama_awq,
+    "gemma3_text": gemma3_text_awq,
+    "gemma3": update(gemma3_text_awq, lm_key="language_model"),
+    "deepseek_v2": deepseek_v2_awq,
+}
+
+
+def mse(x, y):
+    return ((x - y).astype(mx.float32)) ** 2
+
+
+def submodule_from_key(module, key):
+    keys = key.split(".")
+    for k in keys:
+        module = module[k]
+    return module
+
+
+def run_layer(
+    layer: nn.Module,
+    x: mx.array,
+    indices: mx.array | None = None,
+    batch_size: int = 32,
+    **kwargs,
+):
+    y = []
+    for i in range(0, x.shape[0], batch_size):
+        if indices is not None:
+            y.append(
+                layer(x[i : i + batch_size], indices[i : i + batch_size], **kwargs)
+            )
+        else:
+            y.append(layer(x[i : i + batch_size], **kwargs))
+        mx.eval(y)
+    y = mx.concatenate(y, axis=0)
+    return y
+
+
+def dist_split(x: mx.array, group: mx.distributed.Group):
+    N = group.size()
+    if N == 1:
+        return x
+    B = x.shape[0]
+    assert B % N == 0
+    r = group.rank()
+    local_B = (B + N - 1) // N
+    return x[r * local_B : (r + 1) * local_B]
+
+
+def search_best_scale(
+    layers: list[nn.Module],
+    quantize_func: Callable,
+    block: nn.Module | None,
+    layer_kwargs: dict,
+    n_grid: int,
+):
+    group = mx.distributed.init()
+
+    layer_kwargs = layer_kwargs or {}
+
+    x = layers[0].input_feat
+
+    block = block or layers[0]
+    out = block(x, **layer_kwargs)
+
+    x_max = x.abs().mean(axis=(0, 1))
+
+    best_error = float("inf")
+    best_scales = None
+
+    weights = tree_flatten(block.parameters())
+
+    # Search across different scaling ratios
+    # and take the best loss.
+    for ratio in range(n_grid):
+        ratio = ratio / n_grid
+        scales = mx.maximum(x_max**ratio, 1e-4).reshape(-1)
+        scales = scales / (scales.max() * scales.min()).sqrt()
+        for layer in layers:
+            if isinstance(layer, (nn.Linear, SwitchLinear)):
+                layer.weight = quantize_func(layer.weight * scales) / scales
+
+        out_q = run_layer(block, x, **layer_kwargs)
+        loss = mse(out, out_q).sum()
+        if group is not None:
+            loss = mx.distributed.all_sum(loss) / group.size()
+        loss /= out.size
+        mx.eval(loss)
+        if loss.item() < best_error:
+            best_error = loss.item()
+            best_scales = scales
+
+        # reload the original weights
+        block.load_weights(weights)
+
+    best_scales = best_scales.reshape(-1)
+    mx.eval(best_scales)
+    return best_scales
+
+
+def apply_scale(prev_op, layers, scales):
+    # Fuse the scales into the previous op
+    if isinstance(prev_op, (nn.Linear, SwitchLinear)):
+        assert len(layers) == 1
+        prev_op.weight = prev_op.weight / scales[:, mx.newaxis]
+        if hasattr(prev_op, "bias"):
+            prev_op.bias = prev_op.bias / scales
+        layers[0].weight = layers[0].weight * scales
+    elif isinstance(prev_op, (nn.LayerNorm, nn.RMSNorm)):
+        prev_op.weight = prev_op.weight / scales
+        if hasattr(prev_op, "bias"):
+            prev_op.bias = prev_op.bias / scales
+        for layer in layers:
+            layer.weight = layer.weight * scales
+    elif prev_op.__class__.__name__ == "RMSNorm":  # For gemma models
+        dt = prev_op.weight.dtype
+        prev_op.weight = (
+            (1.0 + prev_op.weight.astype(mx.float32)) / scales - 1.0
+        ).astype(dt)
+        for layer in layers:
+            layer.weight = layer.weight * scales
+    else:
+        raise NotImplementedError(f"Could not apply scale to prev_op: {prev_op}")
+
+    for layer in layers:
+        if hasattr(layer, "input_feat"):
+            layer.input_feat = layer.input_feat / scales
+
+
+def scale_block(
+    block: nn.Module,
+    configs: list[ScaleConfig],
+    quantize_func: Callable,
+    layer_kwargs: dict,
+    n_grid: int,
+):
+    for conf in configs:
+        if conf.use_config is not None and not conf.use_config(block):
+            continue
+        if conf.block is not None:
+            local_block = block[conf.block]
+            layers = [submodule_from_key(local_block, l) for l in conf.layers]
+        else:
+            local_block = None
+            layers = [submodule_from_key(block, l) for l in conf.layers]
+        local_kwargs = {k: layer_kwargs[k] for k in conf.kwargs if k in layer_kwargs}
+        for k in conf.kwargs:
+            if hasattr(layers[0], k):
+                local_kwargs[k] = getattr(layers[0], k)
+
+        scales = search_best_scale(
+            layers=layers,
+            block=local_block,
+            layer_kwargs=local_kwargs,
+            quantize_func=quantize_func,
+            n_grid=n_grid,
+        )
+        apply_scale(submodule_from_key(block, conf.prev), layers, scales)
+
+
+def search_best_clip(
+    module: nn.Module,
+    quantize_func: Callable,
+    group_size: int,
+    n_grid: int,
+    max_shrink: float = 0.5,
+    batch_size: int = 64,
+    n_frames: int = 512,
+):
+    group = mx.distributed.init()
+
+    # subsample the input features
+    x = module.input_feat.flatten(0, 1)
+    stride = (x.shape[0] + n_frames - 1) // n_frames
+    x = x[::stride]
+
+    w = module.weight
+    x = x.reshape(x.shape[0], -1, group_size)
+
+    w_init_shape = w.shape
+    w_all = mx.flatten(w, 0, w.ndim - 2)
+    w_max_all = []
+
+    # batch across W to save memory
+    for b in range(0, w_all.shape[0], batch_size):
+        w = w_all[b : b + batch_size]
+
+        group_shape = (w.shape[0], w.shape[-1] // group_size)
+        best_error = mx.full(group_shape, float("inf"))
+        best_w_max = mx.zeros((*group_shape, 1), dtype=x.dtype)
+
+        w_shape = w.shape
+
+        w = w.reshape(*w.shape[:-1], -1, group_size)
+        out = mx.einsum("bdg,odg->bod", x, w)
+        init_max = w.abs().max(axis=-1, keepdims=True)
+
+        # try a range of clips and pick the one with the smallest loss
+        for i in range(int(max_shrink * n_grid)):
+            p = 1 - i / n_grid
+            w_max = p * init_max
+            w_m = mx.clip(w, -w_max, w_max).reshape(w_shape)
+
+            w_q = quantize_func(w_m)
+
+            w_q = w_q.reshape(*w_q.shape[:-1], -1, group_size)
+            out_q = mx.einsum("bdg,odg->bod", x, w_q)
+
+            # Take the mean across the input batch
+            loss = mse(out, out_q).sum(axis=0)
+            if group is not None:
+                loss = mx.distributed.all_sum(loss) / group.size()
+            loss /= out.shape[0]
+            best_indices = loss < best_error
+            best_error = mx.where(best_indices, loss, best_error)
+            best_w_max = mx.where(best_indices[..., mx.newaxis], w_max, best_w_max)
+            mx.eval(best_w_max, best_error)
+
+        w_max_all.append(best_w_max)
+
+    best_w_max = mx.concatenate(w_max_all, axis=0)
+
+    w_r = w_all.reshape(*w_all.shape[:-1], -1, group_size)
+    best_w = mx.clip(w_r, -best_w_max, best_w_max)
+    best_w = best_w.reshape(w_init_shape)
+
+    mx.eval(best_w)
+    return best_w
+
+
+def clip_block(
+    block: nn.Module,
+    no_clip_keys: list[str],
+    quantize_func: Callable,
+    group_size: int,
+    n_grid: int = 20,
+):
+    def apply_clip(path, module):
+        if isinstance(module, (nn.Linear, SwitchLinear)) and all(
+            k not in path for k in no_clip_keys
+        ):
+            best_weight = search_best_clip(
+                module,
+                quantize_func=quantize_func,
+                group_size=group_size,
+                n_grid=n_grid,
+            )
+            module.weight = best_weight
+
+    tree_map_with_path(apply_clip, block.leaf_modules(), is_leaf=nn.Module.is_module)
+
+
+def awq_quantize(
+    model,
+    inputs: mx.array,
+    awq_config: AWQConfig,
+    group_size: int = 64,
+    bits: int = 3,
+    embed_group_size: int = 32,
+    embed_bits: int = 4,
+    n_grid: int = 20,
+):
+    if awq_config.lm_key is not None:
+        model = model[awq_config.lm_key]
+
+    group = mx.distributed.init()
+
+    def quantize_func(w):
+        wq = mx.quantize(w, bits=bits, group_size=group_size)
+        return mx.dequantize(*wq, bits=bits, group_size=group_size)
+
+    mask = create_attention_mask(inputs)
+
+    embed_key = awq_config.embed
+    model.model[embed_key] = model.model[embed_key].to_quantized(
+        group_size=embed_group_size, bits=embed_bits
+    )
+    inputs = model.model[embed_key](inputs)
+
+    def capture(module):
+        if not isinstance(module, (nn.Linear, SwitchLinear)):
+            return module
+
+        class Catcher(nn.Module):
+            def __call__(self, x: mx.array, *args, **kwargs):
+                # Store the input features on the original modules.
+                if hasattr(module, "input_feat"):
+                    module.input_feat = mx.concatenate([module.input_feat, x], axis=0)
+                else:
+                    module.input_feat = x
+
+                # Also store the MOE indices if applicabale
+                if isinstance(module, SwitchLinear):
+                    indices = args[0]
+                    if hasattr(module, "indices"):
+                        module.indices = mx.concatenate(
+                            [module.indices, indices], axis=0
+                        )
+                    else:
+                        module.indices = indices
+
+                return module(x, *args, **kwargs)
+
+        return Catcher()
+
+    for e, block in enumerate(tqdm(model.layers)):
+        # Capture the input features for each of the layers in the transformer block
+        orig_leaves = block.leaf_modules()
+        capture_leaves = tree_map(capture, orig_leaves, is_leaf=nn.Module.is_module)
+        block.update_modules(capture_leaves)
+        outputs = run_layer(block, inputs, mask=mask)
+        block.update_modules(orig_leaves)
+        del capture_leaves
+
+        # Quantize the block without AWQ to obtain a reference loss
+        nn.quantize(block, group_size=group_size, bits=bits)
+        outputs_q = run_layer(block, inputs, mask=mask)
+        before_loss = mse(outputs, outputs_q).sum()
+        if group is not None:
+            before_loss = mx.distributed.all_sum(before_loss) / group.size()
+        before_loss /= outputs.size
+        block.update_modules(orig_leaves)
+        orig_params = block.parameters()
+
+        scale_block(
+            block=block,
+            configs=awq_config.scale_configs,
+            quantize_func=quantize_func,
+            n_grid=n_grid,
+            layer_kwargs={"mask": mask},
+        )
+
+        clip_block(
+            block=block,
+            no_clip_keys=awq_config.no_clip,
+            quantize_func=quantize_func,
+            group_size=group_size,
+            n_grid=n_grid,
+        )
+
+        # Quantize the scaled and clipped block
+        nn.quantize(block, group_size=group_size, bits=bits)
+        outputs_q = run_layer(block, inputs, mask=mask)
+        after_loss = mse(outputs, outputs_q).sum()
+        if group is not None:
+            after_loss = mx.distributed.all_sum(after_loss) / group.size()
+        after_loss /= outputs.size
+        tqdm.write(f"Loss reduction: {after_loss / before_loss}")
+        if after_loss > before_loss:
+            # Reload original weights and quantize
+            block.update_modules(orig_leaves)
+            block.update(orig_params)
+            nn.quantize(block, group_size=group_size, bits=bits)
+            tqdm.write("Loss is not reduced, falling back to original weights.")
+
+        inputs = outputs
+
+        mx.eval(block)
+        mx.clear_cache()
+
+    if (lm_head := awq_config.lm_head) in model:
+        model[lm_head] = model[lm_head].to_quantized(
+            group_size=embed_group_size, bits=embed_bits
+        )
+
+
+def load_dataset(tokenizer, num_samples: int, sequence_length: int) -> mx.array:
+    save_dir = Path.home() / ".cache/mlx-lm/calibration_v5.txt"
+    if not save_dir.exists():
+        save_dir.parent.mkdir(parents=True, exist_ok=True)
+        url = "https://gist.githubusercontent.com/tristandruyen/9e207a95c7d75ddf37525d353e00659c/raw/571fda718462de863e5a0171078c175420c7649a/calibration_data_v5_rc.txt"
+        request.urlretrieve(url, save_dir)
+    with open(save_dir) as fid:
+        texts = fid.read()
+    tokens = tokenizer.encode(texts, return_tensors="mlx")[0]
+
+    # select random non-overlapping chunks
+    tokens = tokens[: (tokens.size // sequence_length) * sequence_length]
+    tokens = tokens.reshape(-1, sequence_length)
+    segments = mx.random.permutation(tokens.shape[0])[:num_samples]
+    return tokens[segments]
+
+
+def update_config(
+    model: nn.Module,
+    config: Dict[str, Any],
+):
+    # dummy
+    config["quantization"] = {"group_size": 64, "bits": 4}
+
+    def update_config(path, module):
+        if hasattr(module, "bits"):
+            config["quantization"][path] = {
+                "group_size": module.group_size,
+                "bits": module.bits,
+            }
+        else:
+            config["quantization"][path] = False
+
+    tree_map_with_path(update_config, model.leaf_modules(), is_leaf=nn.Module.is_module)
+    return config
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model", "-m", default="mlx-community/Qwen2.5-7B-Instruct-bf16"
+    )
+    parser.add_argument("--mlx-path", default="mlx_model")
+    parser.add_argument("--bits", type=int, default=4)
+    parser.add_argument("--group-size", type=int, default=64)
+    parser.add_argument("--embed-bits", type=int, default=4)
+    parser.add_argument("--embed-group-size", type=int, default=32)
+    parser.add_argument("--num-samples", type=int, default=128)
+    parser.add_argument("--sequence-length", type=int, default=512)
+    parser.add_argument("--n-grid", type=int, default=20)
+    parser.add_argument("--seed", type=int, default=123)
+    args = parser.parse_args()
+
+    group = mx.distributed.init()
+
+    num_samples = args.num_samples
+    if group is not None and num_samples % group.size() > 0:
+        num_samples += group.size() - num_samples % group.size()
+
+    mx.random.seed(args.seed)
+
+    model_path = get_model_path(args.model, revision=None)
+    model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
+
+    model_type = config["model_type"]
+    if (awq_config := AWQ_MODEL_CONFIGS.get(model_type, None)) is None:
+        raise NotImplementedError(f"AWQ support for {model_type} models NYI.")
+
+    calibration_data = load_dataset(tokenizer, args.num_samples, args.sequence_length)
+
+    calibration_data = dist_split(calibration_data, group)
+
+    awq_quantize(
+        model,
+        calibration_data,
+        awq_config,
+        bits=args.bits,
+        group_size=args.group_size,
+        embed_bits=args.embed_bits,
+        embed_group_size=args.embed_group_size,
+        n_grid=args.n_grid,
+    )
+
+    config = update_config(model, config)
+    weights = dict(tree_flatten(model.parameters()))
+    save(
+        args.mlx_path,
+        model_path,
+        weights,
+        tokenizer,
+        config,
+        hf_repo=args.model,
+    )

mlx_lm/cache_prompt.py

@@ -148,7 +148,7 @@ def main():
         pass
 
     print()
-    print(f"Peak memory: {mx.metal.get_peak_memory() / 1e9:.3f} GB")
+    print(f"Peak memory: {mx.get_peak_memory() / 1e9:.3f} GB")
 
     print("Saving...")
     metadata = {}
@@ -159,4 +159,8 @@ def main():
 
 
 if __name__ == "__main__":
+    print(
+        "Calling `python -m mlx_lm.cache_prompt...` directly is deprecated."
+        " Use `mlx_lm.cache_prompt...` or `python -m mlx_lm cache_prompt ...` instead."
+    )
     main()

mlx_lm/chat.py

@@ -12,6 +12,8 @@ from .utils import load
 
 DEFAULT_TEMP = 0.0
 DEFAULT_TOP_P = 1.0
+DEFAULT_XTC_PROBABILITY = 0.0
+DEFAULT_XTC_THRESHOLD = 0.0
 DEFAULT_SEED = None
 DEFAULT_MAX_TOKENS = 256
 DEFAULT_MODEL = "mlx-community/Llama-3.2-3B-Instruct-4bit"
@@ -37,6 +39,18 @@ def setup_arg_parser():
     parser.add_argument(
         "--top-p", type=float, default=DEFAULT_TOP_P, help="Sampling top-p"
     )
+    parser.add_argument(
+        "--xtc-probability",
+        type=float,
+        default=DEFAULT_XTC_PROBABILITY,
+        help="Probability of XTC sampling to happen each next token",
+    )
+    parser.add_argument(
+        "--xtc-threshold",
+        type=float,
+        default=0.0,
+        help="Thresold the probs of each next token candidate to be sampled by XTC",
+    )
     parser.add_argument(
         "--seed",
         type=int,
@@ -98,7 +112,15 @@ def main():
             tokenizer,
             prompt,
             max_tokens=args.max_tokens,
-            sampler=make_sampler(args.temp, args.top_p),
+            sampler=make_sampler(
+                args.temp,
+                args.top_p,
+                xtc_threshold=args.xtc_threshold,
+                xtc_probability=args.xtc_probability,
+                xtc_special_tokens=(
+                    tokenizer.encode("\n") + list(tokenizer.eos_token_ids)
+                ),
+            ),
             prompt_cache=prompt_cache,
         ):
             print(response.text, flush=True, end="")
@@ -106,4 +128,8 @@ def main():
 
 
 if __name__ == "__main__":
+    print(
+        "Calling `python -m mlx_lm.chat...` directly is deprecated."
+        " Use `mlx_lm.chat...` or `python -m mlx_lm chat ...` instead."
+    )
     main()

mlx_lm/convert.py

@@ -1,8 +1,6 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import argparse
-import glob
-import shutil
 from pathlib import Path
 from typing import Callable, Optional, Union
 
@@ -15,15 +13,40 @@ from .utils import (
     fetch_from_hub,
     get_model_path,
     quantize_model,
-    save_config,
-    save_weights,
+    save,
     upload_to_hub,
 )
 
 
 def mixed_quant_predicate_builder(
-    low_bits: int = 4, high_bits: int = 4, group_size: int = 64
+    recipe: str, model: nn.Module
 ) -> Callable[[str, nn.Module, dict], Union[bool, dict]]:
+
+    high_bits = 6
+    group_size = 64
+
+    if recipe == "mixed_2_6":
+        low_bits = 2
+    elif recipe == "mixed_3_4":
+        low_bits = 3
+        high_bits = 4
+    elif recipe == "mixed_3_6":
+        low_bits = 3
+    elif recipe == "mixed_4_6":
+        low_bits = 4
+    else:
+        raise ValueError("Invalid quant recipe {recipe}")
+
+    down_keys = [k for k, _ in model.named_modules() if "down_proj" in k]
+    if len(down_keys) == 0:
+        raise ValueError("Model does not have expected keys for mixed quant.")
+
+    # Look for the layer index location in the path:
+    for layer_location, k in enumerate(down_keys[0].split(".")):
+        if k.isdigit():
+            break
+    num_layers = len(model.layers)
+
     def mixed_quant_predicate(
         path: str,
         module: nn.Module,
@@ -37,9 +60,11 @@ def mixed_quant_predicate_builder(
         if not hasattr(module, "to_quantized"):
             return False
 
-        index = int(path.split(".")[2]) if len(path.split(".")) > 2 else 0
-
-        num_layers = config["num_hidden_layers"]
+        index = (
+            int(path.split(".")[layer_location])
+            if len(path.split(".")) > layer_location
+            else 0
+        )
         use_more_bits = (
             index < num_layers // 8
             or index >= 7 * num_layers // 8
@@ -57,19 +82,9 @@ def mixed_quant_predicate_builder(
     return mixed_quant_predicate
 
 
-QUANT_RECIPES = {
-    "mixed_2_6": mixed_quant_predicate_builder(low_bits=3, high_bits=6),
-    "mixed_3_6": mixed_quant_predicate_builder(low_bits=2, high_bits=6),
-}
+QUANT_RECIPES = ["mixed_2_6", "mixed_3_4", "mixed_3_6", "mixed_4_6"]
 
-
-def quant_args(arg):
-    if arg not in QUANT_RECIPES:
-        raise argparse.ArgumentTypeError(
-            f"Invalid q-recipe {arg!r}. Choose from: {list(QUANT_RECIPES.keys())}"
-        )
-    else:
-        return QUANT_RECIPES[arg]
+MODEL_CONVERSION_DTYPES = ["float16", "bfloat16", "float32"]
 
 
 def convert(
@@ -78,12 +93,12 @@ def convert(
     quantize: bool = False,
     q_group_size: int = 64,
     q_bits: int = 4,
-    dtype: str = "float16",
+    dtype: Optional[str] = None,
     upload_repo: str = None,
     revision: Optional[str] = None,
     dequantize: bool = False,
     quant_predicate: Optional[
-        Callable[[str, nn.Module, dict], Union[bool, dict]]
+        Union[Callable[[str, nn.Module, dict], Union[bool, dict]], str]
     ] = None,
 ):
     # Check the save path is empty
@@ -100,9 +115,23 @@ def convert(
     model_path = get_model_path(hf_path, revision=revision)
     model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
 
+    if isinstance(quant_predicate, str):
+        quant_predicate = mixed_quant_predicate_builder(quant_predicate, model)
+
+    if dtype is None:
+        dtype = config.get("torch_dtype", None)
     weights = dict(tree_flatten(model.parameters()))
-    dtype = getattr(mx, dtype)
-    weights = {k: v.astype(dtype) for k, v in weights.items()}
+    if dtype in MODEL_CONVERSION_DTYPES:
+        print("[INFO] Using dtype:", dtype)
+        dtype = getattr(mx, dtype)
+
+        if hasattr(model, "cast_predicate"):
+            cast_predicate = model.cast_predicate()
+        else:
+            cast_predicate = lambda _: True
+        weights = {
+            k: v.astype(dtype) if cast_predicate(k) else v for k, v in weights.items()
+        }
 
     if quantize and dequantize:
         raise ValueError("Choose either quantize or dequantize, not both.")
@@ -120,18 +149,17 @@ def convert(
         weights = dict(tree_flatten(model.parameters()))
 
     del model
-    save_weights(mlx_path, weights, donate_weights=True)
-
-    py_files = glob.glob(str(model_path / "*.py"))
-    for file in py_files:
-        shutil.copy(file, mlx_path)
-
-    tokenizer.save_pretrained(mlx_path)
-
-    save_config(config, config_path=mlx_path / "config.json")
+    save(
+        mlx_path,
+        model_path,
+        weights,
+        tokenizer,
+        config,
+        hf_repo=hf_path,
+    )
 
     if upload_repo is not None:
-        upload_to_hub(mlx_path, upload_repo, hf_path)
+        upload_to_hub(mlx_path, upload_repo)
 
 
 def configure_parser() -> argparse.ArgumentParser:
@@ -160,16 +188,17 @@ def configure_parser() -> argparse.ArgumentParser:
     )
     parser.add_argument(
         "--quant-predicate",
-        help=f"Mixed-bit quantization recipe. Choices: {list(QUANT_RECIPES.keys())}",
-        type=quant_args,
+        help=f"Mixed-bit quantization recipe.",
+        choices=QUANT_RECIPES,
+        type=str,
         required=False,
     )
     parser.add_argument(
         "--dtype",
-        help="Type to save the non-quantized parameters.",
+        help="Type to save the non-quantized parameters. Defaults to config.json's `torch_dtype` or the current model weights dtype.",
         type=str,
-        choices=["float16", "bfloat16", "float32"],
-        default="float16",
+        choices=MODEL_CONVERSION_DTYPES,
+        default=None,
     )
     parser.add_argument(
         "--upload-repo",

mlx_lm/dwq.py

@@ -0,0 +1,254 @@
+# Copyright © 2025 Apple Inc.
+
+import argparse
+import copy
+import glob
+import shutil
+import time
+import types
+from pathlib import Path
+
+import mlx.core as mx
+import mlx.nn as nn
+import mlx.optimizers as optimizers
+import numpy as np
+from mlx.utils import tree_flatten, tree_map
+
+from mlx_lm.tokenizer_utils import TokenizerWrapper
+from mlx_lm.tuner.datasets import load_dataset
+from mlx_lm.tuner.trainer import iterate_batches
+from mlx_lm.tuner.utils import print_trainable_parameters
+from mlx_lm.utils import (
+    create_model_card,
+    fetch_from_hub,
+    get_model_path,
+    quantize_model,
+    save_config,
+    save_weights,
+)
+
+
+class Catcher(nn.Module):
+    def __init__(self, module):
+        super().__init__()
+        self.module = module
+
+    def __call__(self, *args, **kwargs):
+        self.outputs = self.module(*args, **kwargs)
+        return self.outputs
+
+
+def dwq_quantize(
+    model,
+    q_model,
+    opt,
+    data,
+    batch_size: int = 2,
+    max_seq_length: int = 2048,
+    temperature: float = 1.0,
+    activation_layer_step: float = 0.25,
+    activation_loss_weight: float = 1e-1,
+    dtype: mx.Dtype = mx.bfloat16,
+):
+    group = mx.distributed.init()
+    world_size = group.size()
+    rank = group.rank()
+
+    def unfreeze(_, m):
+        if hasattr(m, "bits") and hasattr(m, "group_size"):
+            m.unfreeze(keys=["scales", "biases"], recurse=False)
+
+    q_model.apply_to_modules(unfreeze)
+    print_trainable_parameters(q_model)
+
+    layer_id_step = int(activation_layer_step * len(model.layers))
+    layer_ids = list(range(len(model.layers)))[layer_id_step::layer_id_step]
+
+    for lid in layer_ids:
+        model.layers[lid] = Catcher(model.layers[lid])
+        q_model.layers[lid] = Catcher(q_model.layers[lid])
+
+    def log_norm(x):
+        if temperature != 1.0:
+            x = x * (1 / temperature)
+        return x - mx.logsumexp(x, axis=-1, keepdims=True)
+
+    def forward(model, inputs):
+        logprobs = log_norm(model(inputs).astype(mx.float32))
+        extra_targets = [
+            model.layers[lid].outputs.astype(mx.float32) for lid in layer_ids
+        ]
+        for lid in layer_ids:
+            model.layers[lid].outputs = None
+        return logprobs, extra_targets
+
+    def loss_fn(params, x, targets, extra_targets, lengths):
+        q_model.update(tree_map(lambda x: x.astype(dtype), params))
+        logprobs, q_extra_targets = forward(q_model, x)
+        losses = nn.losses.kl_div_loss(logprobs, targets, reduction="none")
+        mask = mx.arange(targets.shape[1]) < lengths[:, 1:]
+        ntoks = mask.sum()
+        kl_loss = (mask * losses).sum() / ntoks
+        act_loss = mx.stack(
+            [
+                (mask * (qe - e).abs().mean(axis=-1)).sum() / ntoks
+                for qe, e in zip(q_extra_targets, extra_targets)
+            ]
+        )
+        loss = kl_loss + activation_loss_weight * act_loss.mean()
+        return loss, ntoks
+
+    def step(inputs, targets, extra_targets, lengths, params):
+        (loss, ntoks), grads = mx.value_and_grad(loss_fn)(
+            params, inputs, targets, extra_targets, lengths
+        )
+        grads = nn.average_gradients(grads)
+        params = opt.apply_gradients(grads, params)
+        return loss, ntoks, params
+
+    # Accumulate learned weights in higher precision
+    params = tree_map(
+        lambda x: x.astype(mx.float32),
+        q_model.trainable_parameters(),
+    )
+
+    avg_loss = None
+    tokens = 0
+    tic = time.time()
+    for it, (batch, lengths) in enumerate(
+        iterate_batches(data, batch_size, max_seq_length)
+    ):
+        targets, extra_targets = forward(model, batch)
+        mx.eval(targets, extra_targets)
+        loss, ntoks, params = step(batch, targets, extra_targets, lengths, params)
+        mx.eval(loss, params)
+        loss = mx.distributed.all_sum(loss, stream=mx.cpu).item() / world_size
+        ntoks = mx.distributed.all_sum(ntoks, stream=mx.cpu).item()
+        tokens += ntoks
+        toks_per_sec = tokens / (time.time() - tic)
+        avg_loss = 0.95 * (avg_loss or loss) + 0.05 * loss
+        if rank == 0:
+            print(
+                f"{it=}, {loss=:.3f}, {avg_loss=:.4f}, {tokens=}, {toks_per_sec=:.3f}",
+                flush=True,
+            )
+    q_model.update(tree_map(lambda x: x.astype(dtype), params))
+    for lid in layer_ids:
+        q_model.layers[lid] = q_model.layers[lid].module
+
+
+def save_model(
+    model: nn.Module,
+    tokenizer: TokenizerWrapper,
+    config,
+    model_path: Path,
+    mlx_path: str,
+    hf_path: str,
+):
+    weights = dict(tree_flatten(model.parameters()))
+
+    mlx_path = Path(mlx_path)
+    save_weights(mlx_path, weights, donate_weights=True)
+
+    py_files = glob.glob(str(model_path / "*.py"))
+    for file in py_files:
+        shutil.copy(file, mlx_path)
+
+    tokenizer.save_pretrained(mlx_path)
+
+    save_config(config, config_path=mlx_path / "config.json")
+    create_model_card(mlx_path, hf_path)
+
+
+def load_data(tokenizer, data_path: str, num_samples: int):
+    args = types.SimpleNamespace(
+        hf_dataset={
+            "path": data_path,
+            "train_split": f"train",
+            "valid_split": "train[:1]",
+        },
+        train=True,
+        test=False,
+    )
+    dataset = load_dataset(args, tokenizer)[0]
+    perm = np.random.permutation(len(dataset))[:num_samples].tolist()
+    return [dataset.process(dataset[i]) for i in perm]
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", "-m", default="Qwen/Qwen3-4B")
+    parser.add_argument("--quantized-model", default=None)
+    parser.add_argument(
+        "--mlx-path", default="mlx_model", help="Path to save the quantized model."
+    )
+    parser.add_argument(
+        "--bits",
+        type=int,
+        default=4,
+        help="Bits per weight for quantization.",
+    )
+    parser.add_argument(
+        "--group-size", type=int, default=64, help="Group size for quantization."
+    )
+    parser.add_argument(
+        "--num-samples",
+        type=int,
+        default=1024,
+        help="Number of samples to use for training.",
+    )
+    parser.add_argument("--max-seq-length", type=int, default=2048)
+    parser.add_argument("--seed", type=int, default=123)
+    parser.add_argument("--learning-rate", type=float, default=1e-6)
+    parser.add_argument("--batch-size", type=int, default=4)
+    parser.add_argument(
+        "--data-path",
+        type=str,
+        default="allenai/tulu-3-sft-mixture",
+        help="A Hugging Face dataset which is compatible with an mlx-lm dataset format.",
+    )
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=1.0,
+        help="Temperature scaling for the loss.",
+    )
+    args = parser.parse_args()
+
+    group = mx.distributed.init()
+
+    num_samples = args.num_samples
+    if num_samples % group.size() > 0:
+        num_samples += group.size() - num_samples % group.size()
+
+    np.random.seed(args.seed)
+    mx.random.seed(args.seed)
+
+    model_path = get_model_path(args.model, revision=None)
+    model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
+
+    calibration_data = load_data(tokenizer, args.data_path, args.num_samples)
+
+    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)
+    else:
+        q_model = copy.deepcopy(model)
+        _, config = quantize_model(
+            q_model,
+            config,
+            q_group_size=args.group_size,
+            q_bits=args.bits,
+        )
+
+    opt = optimizers.Adam(learning_rate=args.learning_rate, bias_correction=True)
+    dwq_quantize(
+        model,
+        q_model,
+        opt,
+        calibration_data,
+        batch_size=args.batch_size,
+        max_seq_length=args.max_seq_length,
+        temperature=args.temperature,
+    )
+    save_model(q_model, tokenizer, config, model_path, args.mlx_path, args.model)

mlx_lm/evaluate.py

@@ -5,12 +5,14 @@ Adapted from a PyTorch implementation by David Grangier
 """
 
 import argparse
+import collections
+import copy
 import json
 import logging
 import os
 from importlib.metadata import version
 from pathlib import Path
-from typing import Optional, Union
+from typing import Any, Optional
 
 import lm_eval
 import mlx.core as mx
@@ -21,19 +23,9 @@ from lm_eval.api.registry import register_model
 from tqdm import tqdm
 
 from .generate import stream_generate
+from .models.base import create_causal_mask
 from .models.cache import make_prompt_cache
-from .utils import load
-
-PAD = 0
-
-
-def _len_longest_common_prefix(a, b):
-    l = 0
-    for item_a, item_b in zip(a, b):
-        if item_a != item_b:
-            break
-        l += 1
-    return l
+from .utils import common_prefix_len, load
 
 
 def _rstrip_until(s, untils):
@@ -44,72 +36,82 @@ def _rstrip_until(s, untils):
     return s[: min(f)]
 
 
-def _pad_inputs(
-    inputs,
-    maxlen,
-    genlen=0,
-    pad_left=False,
-    pad_multiple=32,
-    truncate=False,
-):
-    # pad the prompts to the left with at least genlen tokens.
-    actual_maxlen = max(len(p) for p in inputs) + genlen
-    if actual_maxlen > maxlen:
-        if not truncate:
-            raise ValueError("Inputs are too long.")
-        else:  # drop begining
-            actual_maxlen = maxlen
-            inputs = [p[max(0, len(p) - maxlen) :] for p in inputs]
-    if pad_multiple > 0:
-        maxlen = (actual_maxlen + pad_multiple - 1) // pad_multiple
-        maxlen *= pad_multiple
-    assert PAD == 0
-    lr = np.array((1, 0) if pad_left else (0, 1))
-    return np.stack(
-        [np.pad(np.array(x, np.int32), lr * (maxlen - len(x))) for x in inputs],
+def _pad_inputs(inputs):
+    lengths = np.array([len(x) for x in inputs])
+    maxlen = lengths.max()
+    padded = np.stack(
+        [np.pad(x, (0, maxlen - len(x))) for x in inputs],
         axis=0,
     )
+    return mx.array(padded), mx.array(lengths)
+
+
+def chat_template_fn(**extra_kwargs):
+    def apply_chat_template(self, chat_history, add_generation_prompt=True) -> str:
+        return self.tokenizer.apply_chat_template(
+            chat_history,
+            tokenize=False,
+            add_generation_prompt=add_generation_prompt,
+            continue_final_message=not add_generation_prompt,
+            **extra_kwargs,
+        )
+
+    return apply_chat_template
 
 
 @register_model("mlxlm")
 class MLXLM(LM):
+
+    tokenizer_name = lm_eval.models.huggingface.HFLM.tokenizer_name
+    apply_chat_template = chat_template_fn()
+
     def __init__(
         self,
         path_or_hf_repo: str,
-        batch_size: int = 16,
         max_tokens: Optional[int] = None,
         use_chat_template: Optional[bool] = None,
     ) -> None:
         super().__init__()
-        self._batch_size = batch_size
         self._model, self.tokenizer = load(path_or_hf_repo)
         self._max_tokens = max_tokens or self.tokenizer.model_max_length
-        self.use_chat_template = use_chat_template or (
-            self.tokenizer.chat_template is not None
-        )
+        self._batch_size = 8
+        self.use_chat_template = use_chat_template
+        if use_chat_template is None:
+            self.use_chat_template = self.tokenizer.chat_template is not None
 
-    def _score_fn(self, inputs, tokenize=True, step_size=32):
-        if tokenize:
-            inputs = self._tokenize(inputs)
-        inputs = _pad_inputs(inputs, self._max_tokens, truncate=False)
-        inputs = mx.array(inputs)
+    def _process_prompt(self, prompt, step_size: int = 2048):
+        prompt = mx.array(prompt)[None]
+        cache = make_prompt_cache(self._model)
+        for i in range(0, prompt.shape[1], step_size):
+            logits = self._model(prompt[:, i : i + step_size], cache=cache)
+            mx.eval([c.state for c in cache])
+            mx.clear_cache()
+        logprobs = nn.log_softmax(logits[:, -1, :].astype(mx.float32))
+        return logprobs, cache
+
+    def _score_fn(self, inputs, cache: Optional[Any] = None, step_size: int = 2048):
+        inputs, lengths = _pad_inputs(inputs)
         inputs, targets = inputs[..., :-1], inputs[..., 1:]
 
-        cache = make_prompt_cache(self._model)
-
-        mask = targets != PAD
+        cache = cache or make_prompt_cache(self._model)
+        lengths += cache[0].offset
 
         scores, is_greedy = [], []
         for i in range(0, inputs.shape[1], step_size):
-            logits = self._model(inputs[:, i : i + step_size], cache=cache)
+            inp = inputs[:, i : i + step_size]
+            T = inp.shape[1]
 
+            offset = cache[0].offset
+            mask = create_causal_mask(T, offset, lengths=lengths)
+
+            logits = self._model(inp, cache=cache, mask=mask)
             log_probs = nn.log_softmax(logits.astype(mx.float32))
+
             score = mx.take_along_axis(
                 log_probs, targets[:, i : i + step_size, mx.newaxis], axis=-1
             )[..., 0]
-            ig = mask[:, i : i + step_size] * (
-                targets[:, i : i + step_size] == mx.argmax(logits, axis=-1)
-            )
+            ig = targets[:, i : i + step_size] == mx.argmax(logits, axis=-1)
+            ig = mx.where(mx.arange(T) + offset < lengths[:, None], ig, False)
 
             mx.eval(score, ig)
             mx.clear_cache()
@@ -120,38 +122,7 @@ class MLXLM(LM):
         scores = mx.concatenate(scores, axis=1)
         is_greedy = mx.concatenate(is_greedy, axis=1)
 
-        return scores, mask.sum(axis=-1), is_greedy
-
-    def _loglikelihood(self, texts, score_spans=None, tokenize=True):
-        # sort by length to get batches with little padding.
-        sorted_indices = sorted(range(len(texts)), key=lambda i: -len(texts[i]))
-        sorted_inputs = [texts[sorted_indices[i]] for i in range(len(texts))]
-        sorted_spans = None
-        if score_spans is not None:
-            sorted_spans = [score_spans[sorted_indices[i]] for i in range(len(texts))]
-
-        results = []
-        for i in tqdm(range(0, len(sorted_inputs), self._batch_size)):
-            batch = sorted_inputs[i : i + self._batch_size]
-            scores, length, is_greedy = self._score_fn(batch, tokenize=tokenize)
-            for j in range(len(batch)):
-                if sorted_spans is None:  # full sequence score
-                    mask = mx.arange(scores[j].shape[-1]) < length
-                    score = (scores[j].astype(mx.float32) * mask).sum(axis=-1)
-                    ig = (is_greedy[j].astype(mx.int32) * mask).sum(axis=-1)
-                else:  # subsequence score
-                    start, end = sorted_spans[i + j]
-                    score = scores[j][start:end].astype(mx.float32).sum()
-                    ig = is_greedy[j][start:end].astype(mx.int32).sum()
-                    length = end - start
-
-                results.append((score.item(), ig.item(), length))
-
-        # reorder the outputs
-        inv_sort = np.argsort(sorted_indices)
-        results = [results[inv_sort[i]] for i in range(len(results))]
-
-        return results
+        return scores, lengths, is_greedy
 
     def _tokenize(self, texts):
         return [
@@ -183,39 +154,65 @@ class MLXLM(LM):
         """
         logging.info("Estimating loglikelihood for %d pairs." % len(requests))
 
-        # tokenize prefix and prefix + completion for all requests.
-        tokenized = self._tokenize(
-            [t for r in requests for t in [r.args[0], r.args[0] + r.args[1]]]
-        )
+        group = mx.distributed.init()
 
-        # max length (prefix + completion) and longest common prefix per question.
-        length_stats = {}
-        for prefix, completed in zip(tokenized[0::2], tokenized[1::2]):
-            max_completed_l, min_prefix_l = length_stats.get(prefix, (0, 1e8))
-            length_stats[prefix] = (
-                max(max_completed_l, len(completed)),
-                min(min_prefix_l, _len_longest_common_prefix(prefix, completed)),
-            )
+        # Group by common prefix
+        group_reqs = collections.defaultdict(list)
+        for idx, req in enumerate(requests):
+            group_reqs[req.args[0]].append((idx, req.args[1]))
+        questions = list(group_reqs.keys())
+        responses = []
+        indices = []
+        for v in group_reqs.values():
+            idx, resp = zip(*v)
+            indices.extend(idx)
+            responses.append(resp)
+
+        # split data accross ranks
+        questions = questions[group.rank() :: group.size()]
+        responses = responses[group.rank() :: group.size()]
 
-        # truncate requests for completed sequences longer than model context.
-        shortened = []
-        completion_spans = []
         long_completions = 0
-        for prefix, completed in zip(tokenized[0::2], tokenized[1::2]):
-            max_completed_l, prefix_l = length_stats[prefix]
+        scores, is_greedy = [], []
+        for q, rs in tqdm(zip(questions, responses), total=len(questions)):
+            prefix = self._tokenize([q])[0]
+            full_sequences = self._tokenize([q + r for r in rs])
+            max_completed_l = max(len(s) for s in full_sequences)
+
             # compute truncation length
             truncation = max(0, max_completed_l - self._max_tokens - 1)
-            prefix_l = prefix_l - truncation
-            if prefix_l <= 0:
-                # completion too long, prefix is eliminated for some requests.
+            orig_prefix_l = len(prefix)
+            prefix_l = max(len(prefix) - truncation, 0)
+            prefix = prefix[len(prefix) - prefix_l :]
+
+            # If the entire prompt got truncated ignore the question
+            if prefix_l == 0:
                 long_completions += 1
-                truncation = max(0, len(completed) - self._max_tokens - 1)
-                prefix_l = 1
-            # truncate the completed sequence
-            completed = completed[truncation:]
-            shortened.append(completed)
-            # scores do not include initial bos, substract 1 to span bounds
-            completion_spans.append((prefix_l - 1, len(completed) - 1))
+                all_scores.extend([-float("inf")] * len(rs))
+                all_is_greedy.extend([False] * len(rs))
+                continue
+
+            # model scoring, returns num_requests x (logp, is_greedy, length).
+            logprobs, cache = self._process_prompt(prefix)
+            max_idx = mx.argmax(logprobs).item()
+
+            for s in full_sequences:
+                inputs = s[len(prefix) :]
+                # The logprobs from the last token of the prompt are
+                # for the first input token
+                scores.append(logprobs[0, inputs[0]].item())
+                is_greedy.append((inputs[0] == max_idx))
+
+                if len(inputs) == 1:
+                    continue
+                score, _, ig = self._score_fn(
+                    mx.array(inputs)[None, :], cache=copy.deepcopy(cache)
+                )
+                scores[-1] += mx.sum(score).item()
+                is_greedy[-1] &= mx.all(ig).item()
+
+        scores = mx.array(scores)
+        is_greedy = mx.array(is_greedy)
 
         if long_completions > 0:
             logging.info(
@@ -223,16 +220,23 @@ class MLXLM(LM):
                 + "completion longer than context."
             )
 
-        # model scoring, returns num_requests x (logp, is_greedy, length).
-        results = self._loglikelihood(
-            shortened,
-            score_spans=completion_spans,
-            tokenize=False,
-        )
-        return [(r[0], r[1] == r[2]) for r in results]
+        num_results = len(requests)
 
-    tokenizer_name = lm_eval.models.huggingface.HFLM.tokenizer_name
-    apply_chat_template = lm_eval.models.huggingface.HFLM.apply_chat_template
+        # all gather the results across groups
+        if group.size() > 1:
+            per_group = int(np.ceil(num_results / group.size()))
+            scores = mx.pad(scores, ((0, per_group - len(scores)),))
+            is_greedy = mx.pad(is_greedy, ((0, per_group - len(is_greedy))))
+            scores = mx.distributed.all_gather(scores[mx.newaxis], stream=mx.cpu)
+            is_greedy = mx.distributed.all_gather(is_greedy[mx.newaxis], stream=mx.cpu)
+            mx.eval(scores, is_greedy)
+            scores = scores.T.reshape(-1)
+            is_greedy = is_greedy.T.reshape(-1)
+
+        inv_sort = mx.argsort(mx.array(indices))
+        scores = scores[:num_results][inv_sort]
+        is_greedy = is_greedy[:num_results][inv_sort]
+        return list(zip(scores.tolist(), is_greedy.tolist()))
 
     def loglikelihood_rolling(self, requests) -> list[float]:
         """Compute full log-likelihood of a string, with no truncation, for perplexity computation
@@ -269,8 +273,15 @@ class MLXLM(LM):
         logging.info(
             "Estimating loglikelihood rolling for %d sequences." % len(requests)
         )
-        inputs = [req.args[0] for req in requests]
-        return [t[0] for t in self._loglikelihood(inputs)]
+        inputs = self._tokenize([req.args[0] for req in requests])
+        all_scores = []
+        for i in tqdm(range(0, len(texts), self._batch_size)):
+            batch = texts[i : i + self._batch_size]
+            scores, lengths, _ = self._score_fn(batch)
+            mask = mx.arange(scores.shape[-1]) < lengths[:, None]
+            all_scores.extend((mask * scores).sum(axis=-1).tolist())
+
+        return all_scores
 
     def generate_until(self, requests) -> list[str]:
         """Generate greedily until a stopping sequence
@@ -325,7 +336,7 @@ def main():
         "--output-dir", default=".", help="Output directory for result files."
     )
     parser.add_argument("--batch-size", type=int, default=16, help="Batch size")
-    parser.add_argument("--num-shots", type=int, default=0, help="Number of shots")
+    parser.add_argument("--num-shots", type=int, default=None, help="Number of shots")
     parser.add_argument(
         "--max-tokens",
         type=int,
@@ -333,7 +344,7 @@ def main():
     )
     parser.add_argument(
         "--limit",
-        default=100,
+        default=None,
         help="Limit the number of examples per task.",
         type=int,
     )
@@ -353,6 +364,14 @@ def main():
         "otherwise `False`.",
         default=None,
     )
+    parser.add_argument(
+        "--chat-template-args",
+        type=json.loads,
+        help="""A JSON formatted string of arguments for the tokenizer's "
+        "apply_chat_template, e.g. '{"enable_thinking":false}'""",
+        default="{}",
+    )
+
     args = parser.parse_args()
 
     output_dir = Path(args.output_dir)
@@ -365,10 +384,11 @@ def main():
 
     lm = MLXLM(
         args.model,
-        batch_size=args.batch_size,
         max_tokens=args.max_tokens,
         use_chat_template=args.apply_chat_template,
     )
+    MLXLM.apply_chat_template = chat_template_fn(**args.chat_template_args)
+
     results = lm_eval.simple_evaluate(
         model=lm,
         tasks=args.tasks,
@@ -382,12 +402,14 @@ def main():
         fewshot_random_seed=args.seed,
     )
 
-    model_name = args.model.replace("/", "_")
-    task_names = "_".join(args.tasks)
-    ver = version("lm_eval")
-    filename = f"eval_{model_name}_{task_names}_{args.num_shots:02d}_v_{ver}.json"
-    output_path = output_dir / filename
-    output_path.write_text(json.dumps(results["results"], indent=4))
-    print("Results:")
-    for result in results["results"].values():
-        print(json.dumps(result, indent=4))
+    file_keys = ["eval", args.model.replace("/", "_"), version("lm_eval")]
+    if args.num_shots is not None:
+        file_keys += [f"{args.num_shots:02d}"]
+    file_keys += args.tasks
+    filename = "_".join(file_keys)
+    if mx.distributed.init().rank() == 0:
+        output_path = output_dir / filename
+        output_path.write_text(json.dumps(results["results"], indent=4))
+        print("Results:")
+        for result in results["results"].values():
+            print(json.dumps(result, indent=4))

mlx_lm/examples/pipeline_generate.py

@@ -5,8 +5,7 @@ Run with:
 
 ```
 mlx.launch \
- --hostfile /path/to/hosts.txt \
- --backend mpi \
+ --hostfile /path/to/hosts.json \
  /path/to/pipeline_generate.py \
  --prompt "hello world"
 ```
@@ -19,6 +18,7 @@ https://ml-explore.github.io/mlx/build/html/usage/distributed.html).
 
 import argparse
 import json
+import resource
 from pathlib import Path
 
 import mlx.core as mx
@@ -28,6 +28,9 @@ from mlx.utils import tree_flatten
 from mlx_lm import load, stream_generate
 from mlx_lm.utils import load_model, load_tokenizer
 
+# Needed for 8 bit model
+resource.setrlimit(resource.RLIMIT_NOFILE, (2048, 4096))
+
 
 def download(repo: str, allow_patterns: list[str]) -> Path:
     return Path(
@@ -49,7 +52,7 @@ def shard_and_load(repo):
     # which weights we need
     model, _ = load_model(model_path, lazy=True, strict=False)
 
-    group = mx.distributed.init(backend="mpi")
+    group = mx.distributed.init()
     rank = group.rank()
     model.model.pipeline(group)
 
@@ -98,7 +101,7 @@ if __name__ == "__main__":
     )
     args = parser.parse_args()
 
-    group = mx.distributed.init(backend="mpi")
+    group = mx.distributed.init()
     rank = group.rank()
 
     def rprint(*args, **kwargs):

mlx_lm/fuse.py

@@ -1,6 +1,4 @@
 import argparse
-import glob
-import shutil
 from pathlib import Path
 
 from mlx.utils import tree_flatten, tree_unflatten
@@ -12,8 +10,7 @@ from .tuner.utils import dequantize, load_adapters
 from .utils import (
     fetch_from_hub,
     get_model_path,
-    save_config,
-    save_weights,
+    save,
     upload_to_hub,
 )
 
@@ -89,23 +86,21 @@ def main() -> None:
     if args.de_quantize:
         print("De-quantizing model")
         model = dequantize(model)
+        config.pop("quantization", None)
 
     weights = dict(tree_flatten(model.parameters()))
 
     save_path = Path(args.save_path)
-
-    save_weights(save_path, weights)
-
-    py_files = glob.glob(str(model_path / "*.py"))
-    for file in py_files:
-        shutil.copy(file, save_path)
-
-    tokenizer.save_pretrained(save_path)
-
-    if args.de_quantize:
-        config.pop("quantization", None)
-
-    save_config(config, config_path=save_path / "config.json")
+    hf_path = args.hf_path or (args.model if not Path(args.model).exists() else None)
+    save(
+        save_path,
+        model_path,
+        weights,
+        tokenizer,
+        config,
+        hf_repo=hf_path,
+        donate_weights=False,
+    )
 
     if args.export_gguf:
         model_type = config["model_type"]
@@ -116,9 +111,6 @@ def main() -> None:
         convert_to_gguf(model_path, weights, config, str(save_path / args.gguf_path))
 
     if args.upload_repo is not None:
-        hf_path = args.hf_path or (
-            args.model if not Path(args.model).exists() else None
-        )
         if hf_path is None:
             raise ValueError(
                 "Must provide original Hugging Face repo to upload local model."
@@ -127,4 +119,8 @@ def main() -> None:
 
 
 if __name__ == "__main__":
+    print(
+        "Calling `python -m mlx_lm.fuse...` directly is deprecated."
+        " Use `mlx_lm.fuse...` or `python -m mlx_lm fuse ...` instead."
+    )
     main()

mlx_lm/generate.py

@@ -38,6 +38,9 @@ DEFAULT_MAX_TOKENS = 100
 DEFAULT_TEMP = 0.0
 DEFAULT_TOP_P = 1.0
 DEFAULT_MIN_P = 0.0
+DEFAULT_TOP_K = 0
+DEFAULT_XTC_PROBABILITY = 0.0
+DEFAULT_XTC_THRESHOLD = 0.0
 DEFAULT_MIN_TOKENS_TO_KEEP = 1
 DEFAULT_SEED = None
 DEFAULT_MODEL = "mlx-community/Llama-3.2-3B-Instruct-4bit"
@@ -104,6 +107,21 @@ def setup_arg_parser():
     parser.add_argument(
         "--min-p", type=float, default=DEFAULT_MIN_P, help="Sampling min-p"
     )
+    parser.add_argument(
+        "--top-k", type=int, default=DEFAULT_TOP_K, help="Sampling top-k"
+    )
+    parser.add_argument(
+        "--xtc-probability",
+        type=float,
+        default=DEFAULT_XTC_PROBABILITY,
+        help="Probability of XTC sampling to happen each next token",
+    )
+    parser.add_argument(
+        "--xtc-threshold",
+        type=float,
+        default=0.0,
+        help="Thresold the probs of each next token candidate to be sampled by XTC",
+    )
     parser.add_argument(
         "--min-tokens-to-keep",
         type=int,
@@ -806,7 +824,16 @@ def main():
             raise ValueError("Draft model tokenizer does not match model tokenizer.")
     else:
         draft_model = None
-    sampler = make_sampler(args.temp, args.top_p, args.min_p, args.min_tokens_to_keep)
+    sampler = make_sampler(
+        args.temp,
+        args.top_p,
+        args.min_p,
+        args.min_tokens_to_keep,
+        top_k=args.top_k,
+        xtc_probability=args.xtc_probability,
+        xtc_threshold=args.xtc_threshold,
+        xtc_special_tokens=tokenizer.encode("\n") + list(tokenizer.eos_token_ids),
+    )
     response = generate(
         model,
         tokenizer,

mlx_lm/lora.py

@@ -13,8 +13,7 @@ import mlx.optimizers as optim
 import numpy as np
 import yaml
 
-from .tokenizer_utils import TokenizerWrapper
-from .tuner.datasets import load_dataset
+from .tuner.datasets import CacheDataset, load_dataset
 from .tuner.trainer import TrainingArgs, TrainingCallback, evaluate, train
 from .tuner.utils import (
     build_schedule,
@@ -67,7 +66,7 @@ CONFIG_DEFAULTS = {
     "config": None,
     "grad_checkpoint": False,
     "lr_schedule": None,
-    "lora_parameters": {"rank": 8, "alpha": 16, "dropout": 0.0, "scale": 10.0},
+    "lora_parameters": {"rank": 8, "dropout": 0.0, "scale": 10.0},
     "mask_prompt": False,
 }
 
@@ -187,7 +186,6 @@ def build_parser():
 def train_model(
     args,
     model: nn.Module,
-    tokenizer: TokenizerWrapper,
     train_set,
     valid_set,
     training_callback: TrainingCallback = None,
@@ -258,20 +256,18 @@ def train_model(
     # Train model
     train(
         model=model,
-        tokenizer=tokenizer,
         args=training_args,
         optimizer=opt,
-        train_dataset=train_set,
-        val_dataset=valid_set,
+        train_dataset=CacheDataset(train_set),
+        val_dataset=CacheDataset(valid_set),
         training_callback=training_callback,
     )
 
 
-def evaluate_model(args, model: nn.Module, tokenizer: TokenizerWrapper, test_set):
+def evaluate_model(args, model: nn.Module, test_set):
     test_loss = evaluate(
         model=model,
-        dataset=test_set,
-        tokenizer=tokenizer,
+        dataset=CacheDataset(test_set),
         batch_size=args.batch_size,
         num_batches=args.test_batches,
         max_seq_length=args.max_seq_length,
@@ -298,13 +294,13 @@ def run(args, training_callback: TrainingCallback = None):
 
     elif args.train:
         print("Training")
-        train_model(args, model, tokenizer, train_set, valid_set, training_callback)
+        train_model(args, model, train_set, valid_set, training_callback)
     else:
         raise ValueError("Must provide at least one of --train or --test")
 
     if args.test:
         print("Testing")
-        evaluate_model(args, model, tokenizer, test_set)
+        evaluate_model(args, model, test_set)
 
 
 def main():
@@ -330,4 +326,8 @@ def main():
 
 
 if __name__ == "__main__":
+    print(
+        "Calling `python -m mlx_lm.lora...` directly is deprecated."
+        " Use `mlx_lm.lora...` or `python -m mlx_lm lora ...` instead."
+    )
     main()

mlx_lm/manage.py

@@ -136,4 +136,8 @@ def main():
 
 
 if __name__ == "__main__":
+    print(
+        "Calling `python -m mlx_lm.manage...` directly is deprecated."
+        " Use `mlx_lm.manage...` or `python -m mlx_lm manage ...` instead."
+    )
     main()

mlx_lm/merge.py

@@ -169,4 +169,8 @@ def main():
 
 
 if __name__ == "__main__":
+    print(
+        "Calling `python -m mlx_lm.merge...` directly is deprecated."
+        " Use `mlx_lm.merge...` or `python -m mlx_lm merge ...` instead."
+    )
     main()

mlx_lm/models/baichuan_m1.py

@@ -0,0 +1,226 @@
+# 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 CacheList, KVCache, MambaCache, RotatingKVCache
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    rope_theta: float
+    sliding_window: int
+    sliding_window_layers: List[int]
+    conv_window: int
+    rms_norm_eps: float
+    model_type: str = "baichuan_m1"
+    num_swa_attention_heads: Optional[int] = None
+    num_swa_key_value_heads: Optional[int] = None
+    tie_word_embeddings: bool = False
+
+
+class Attention(nn.Module):
+    def __init__(self, config: ModelArgs, layer_idx: Optional[int] = None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        if layer_idx is None:
+            raise ValueError("Layer index must be provided to Attention module.")
+
+        self.is_swa = layer_idx in config.sliding_window_layers
+        self.num_heads = (
+            config.num_swa_attention_heads
+            if self.is_swa and config.num_swa_attention_heads
+            else config.num_attention_heads
+        )
+        self.num_kv_heads = (
+            config.num_swa_key_value_heads
+            if self.is_swa and config.num_swa_key_value_heads
+            else config.num_key_value_heads
+        )
+
+        self.hidden_size = config.hidden_size
+        self.head_dim = self.hidden_size // self.num_heads
+        assert self.head_dim * self.num_heads == self.hidden_size
+
+        self.scale = self.head_dim**-0.5
+
+        self.W_pack = nn.Linear(
+            config.hidden_size,
+            self.hidden_size + 2 * self.num_kv_heads * self.head_dim,
+            bias=False,
+        )
+        self.o_proj = nn.Linear(
+            self.num_heads * self.head_dim, config.hidden_size, bias=False
+        )
+
+        self.rope = nn.RoPE(self.head_dim, traditional=False, base=config.rope_theta)
+
+        self.conv_window = config.conv_window
+        assert self.conv_window == 2
+        self.conv_k = mx.zeros((1, 1, self.num_kv_heads, 1, self.conv_window))
+        self.conv_v = mx.zeros((1, 1, self.num_kv_heads, 1, self.conv_window))
+
+    def _custom_convolution(self, u, weights, state=None):
+        B, H, L, D = u.shape
+        weights = weights.reshape((1, H, self.conv_window, 1, 1))
+        w0 = weights[:, :, 0]
+        w1 = weights[:, :, 1]
+        if state is None:
+            state = mx.zeros((B, H, 1, D), u.dtype)
+        if L > 1:
+            u_prev = mx.concatenate([state, u[:, :, :-1]], axis=2)
+        else:
+            u_prev = state
+        return u_prev * w0 + u * w1
+
+    def __call__(
+        self, x: mx.array, mask: mx.array = None, cache: Any = None
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        proj = self.W_pack(x)
+        q, k, v = mx.split(proj, (D, D + self.num_kv_heads * self.head_dim), axis=-1)
+
+        q = q.reshape(B, L, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
+        k = k.reshape(B, L, self.num_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
+        v = v.reshape(B, L, self.num_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            offset = cache[1].offset
+            last_k, last_v = cache[0][0], cache[0][1]
+        else:
+            offset = 0
+            last_k, last_v = None, None
+
+        k_init = k
+        v_init = v
+        k = self._custom_convolution(k, self.conv_k, state=last_k)
+        v = self._custom_convolution(v, self.conv_v, state=last_v)
+        q = self.rope(q, offset=offset)
+        k = self.rope(k, offset=offset)
+
+        if cache is not None:
+            k, v = cache[1].update_and_fetch(k, v)
+            if L > 0:
+                cache[0][0] = k_init[:, :, -1:, :]
+                cache[0][1] = v_init[:, :, -1:, :]
+
+        out = scaled_dot_product_attention(
+            q, k, v, cache=cache[1], scale=self.scale, mask=mask
+        )
+        out = out.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(out)
+
+
+class MLP(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.gate_proj = nn.Linear(
+            config.hidden_size, config.intermediate_size, bias=False
+        )
+        self.up_proj = nn.Linear(
+            config.hidden_size, config.intermediate_size, bias=False
+        )
+        self.down_proj = nn.Linear(
+            config.intermediate_size, config.hidden_size, bias=False
+        )
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, config: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.self_attn = Attention(config, layer_idx)
+        self.mlp = 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: mx.array = None, cache: Any = None
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        x = x + r
+        r = self.mlp(self.post_attention_layernorm(x))
+        return x + r
+
+
+class BaichuanModel(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.args = config
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = [DecoderLayer(config, i) for i in range(config.num_hidden_layers)]
+        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def __call__(
+        self, inputs: mx.array, mask: mx.array = None, cache: Any = None
+    ) -> mx.array:
+        x = self.embed_tokens(inputs)
+        if mask is None:
+            if cache is not None:
+                c = [cache[0][1]]
+            mask = create_attention_mask(x, c)
+        if cache is None:
+            cache = [None] * len(self.layers)
+        for layer, c in zip(self.layers, cache):
+            x = layer(x, mask, c)
+        return self.norm(x)
+
+
+class Model(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.model_type = config.model_type
+        self.model = BaichuanModel(config)
+        self.tie_word_embeddings = config.tie_word_embeddings
+        if not config.tie_word_embeddings:
+            self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+    def make_cache(self) -> List[Any]:
+        caches = []
+        for i, layer in enumerate(self.model.layers):
+            is_swa = i in self.config.sliding_window_layers
+            conv_cache = MambaCache()
+            if is_swa:
+                kv_cache = RotatingKVCache(max_size=self.config.sliding_window)
+            else:
+                kv_cache = KVCache()
+            caches.append(CacheList(conv_cache, kv_cache))
+        return caches
+
+    def sanitize(self, weights: dict) -> dict:
+        is_quantized = "lm_head.scales" in weights
+        if not is_quantized and "lm_head.weight" in weights:
+            w = weights["lm_head.weight"]
+            dtype = w.dtype
+            w = w.astype(mx.float32)
+            norm = mx.linalg.norm(w, axis=-1, keepdims=True)
+            w = (w / (norm + 1e-7)).astype(dtype)
+            weights["lm_head.weight"] = w
+        return weights
+
+    def __call__(
+        self, inputs: mx.array, mask: mx.array = None, cache: Any = None
+    ) -> mx.array:
+        outputs = self.model(inputs, mask, cache)
+        return self.lm_head(outputs)
+
+    @property
+    def layers(self) -> List[nn.Module]:
+        return self.model.layers

mlx_lm/models/base.py

@@ -89,7 +89,15 @@ def quantized_scaled_dot_product_attention(
         queries, *q_keys, transpose=True, group_size=group_size, bits=bits
     )
     if mask is not None:
-        scores += mask
+        if isinstance(mask, str):
+            qL, kL = scores.shape[-2:]
+            q_indices = mx.arange(kL - qL, kL)
+            k_indices = mx.arange(kL)
+            mask = q_indices[:, None] >= k_indices[None]
+        if mask.dtype == mx.bool_:
+            scores = mx.where(mask, scores, mx.finfo(scores.dtype).min)
+        else:
+            scores += mask
     scores = mx.softmax(scores, axis=-1, precise=True)
     out = mx.quantized_matmul(
         scores, *q_values, transpose=False, group_size=group_size, bits=bits

mlx_lm/models/cache.py

@@ -436,3 +436,76 @@ class MambaCache(_BaseCache):
     @state.setter
     def state(self, v):
         self.cache = v
+
+
+class ChunkedKVCache(KVCache):
+    def __init__(self, chunk_size=None):
+        super().__init__()
+        self.chunk_size = chunk_size
+        self.start_position = 0
+
+    def maybe_trim_front(self):
+        # Maintain the cache below the chunk size
+        if self.keys is not None and self.keys.shape[2] >= self.chunk_size:
+            self.start_position += self.keys.shape[2] - self.chunk_size
+            self.keys = self.keys[..., -self.chunk_size :, :]
+            self.values = self.values[..., -self.chunk_size :, :]
+
+    def update_and_fetch(self, keys, values):
+        prev = self.offset - self.start_position
+        if self.keys is None or (prev + keys.shape[2]) > self.keys.shape[2]:
+            B, n_kv_heads, _, k_head_dim = keys.shape
+            v_head_dim = values.shape[3]
+            n_steps = (self.step + keys.shape[2] - 1) // self.step
+            k_shape = (B, n_kv_heads, n_steps * self.step, k_head_dim)
+            v_shape = (B, n_kv_heads, n_steps * self.step, v_head_dim)
+            new_k = mx.zeros(k_shape, keys.dtype)
+            new_v = mx.zeros(v_shape, values.dtype)
+            if self.keys is not None:
+                if prev % self.step != 0:
+                    self.keys = self.keys[..., :prev, :]
+                    self.values = self.values[..., :prev, :]
+                self.keys = mx.concatenate([self.keys, new_k], axis=2)
+                self.values = mx.concatenate([self.values, new_v], axis=2)
+            else:
+                self.keys, self.values = new_k, new_v
+
+        self.offset += keys.shape[2]
+        end = self.offset - self.start_position
+        self.keys[..., prev:end, :] = keys
+        self.values[..., prev:end, :] = values
+        return self.keys[..., :end, :], self.values[..., :end, :]
+
+    def trim(self, n):
+        n = min(self.offset - self.start_position, n)
+        self.offset -= n
+        return n
+
+    @property
+    def meta_state(self):
+        return tuple(map(str, (self.chunk_size, self.start_position)))
+
+    @meta_state.setter
+    def meta_state(self, v):
+        self.chunk_size, self.start_position = map(int, v)
+
+
+class CacheList(KVCache):
+    def __init__(self, *caches):
+        self.caches = caches
+
+    def __getitem__(self, idx):
+        return self.caches[idx]
+
+    @property
+    def state(self):
+        return [s for c in self.caches for s in c.state]
+
+    @state.setter
+    def state(self, v):
+        state_lens = [len(c.state) for c in self.caches]
+        start = 0
+        for c in self.caches:
+            l = len(c.state)
+            c.state = v[start : start + l]
+            start += l

mlx_lm/models/deepseek_v2.py

@@ -148,7 +148,7 @@ class DeepseekV2Attention(nn.Module):
             self.q_a_proj = nn.Linear(
                 self.hidden_size, self.q_lora_rank, bias=config.attention_bias
             )
-            self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank)
+            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
             )
@@ -158,7 +158,7 @@ class DeepseekV2Attention(nn.Module):
             self.kv_lora_rank + self.qk_rope_head_dim,
             bias=config.attention_bias,
         )
-        self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank)
+        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
@@ -400,8 +400,6 @@ class DeepseekV2Model(nn.Module):
 
         pipeline_rank = self.pipeline_rank
         pipeline_size = self.pipeline_size
-        # Hack to avoid time-outs during prompt-processing
-        dist_stream = mx.cpu if h.shape[1] > 1 else mx.gpu
         if mask is None:
             mask = create_attention_mask(h, cache)
 
@@ -410,19 +408,17 @@ class DeepseekV2Model(nn.Module):
 
         # Receive from the previous process in the pipeline
         if pipeline_rank < pipeline_size - 1:
-            h = mx.distributed.recv_like(h, (pipeline_rank + 1), stream=dist_stream)
+            h = mx.distributed.recv_like(h, (pipeline_rank + 1))
 
         for i in range(self.num_layers):
             h = self.layers[self.start_idx + i](h, mask, cache[i])
 
         # Send to the next process in the pipeline
         if pipeline_rank != 0:
-            h = mx.distributed.send(
-                h, (pipeline_rank - 1) % pipeline_size, stream=dist_stream
-            )
+            h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
 
         # Broadcast h while keeping it in the graph
-        h = mx.distributed.all_gather(h, stream=dist_stream)[: h.shape[0]]
+        h = mx.distributed.all_gather(h)[: h.shape[0]]
 
         return self.norm(h)
 

mlx_lm/models/deepseek_v3.py

@@ -97,9 +97,7 @@ class DeepseekV3YarnRotaryEmbedding(nn.Module):
             scaling_factor, mscale_all_dim
         )
         freq_extra = base ** (mx.arange(0, dim, 2, dtype=mx.float32) / dim)
-        freq_inter = scaling_factor * base ** (
-            mx.arange(0, dim, 2, dtype=mx.float32) / dim
-        )
+        freq_inter = scaling_factor * freq_extra
         low, high = yarn_find_correction_range(
             beta_fast,
             beta_slow,
@@ -157,7 +155,7 @@ class DeepseekV3Attention(nn.Module):
             self.q_a_proj = nn.Linear(
                 self.hidden_size, self.q_lora_rank, bias=config.attention_bias
             )
-            self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank)
+            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
             )
@@ -167,7 +165,7 @@ class DeepseekV3Attention(nn.Module):
             self.kv_lora_rank + self.qk_rope_head_dim,
             bias=config.attention_bias,
         )
-        self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank)
+        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
@@ -291,6 +289,7 @@ def group_expert_select(
 
     k = top_k
     scores = mx.sigmoid(gates.astype(mx.float32))
+    orig_scores = scores
     scores = scores + e_score_correction_bias
     scores = mx.unflatten(scores, axis=-1, shape=(n_group, -1))
     group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True)
@@ -301,9 +300,9 @@ def group_expert_select(
 
     k = top_k
     inds = mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k]
-    scores = mx.take_along_axis(scores, inds, axis=-1)
+    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) + 1e-20
+        denominator = scores.sum(axis=-1, keepdims=True)
         scores = scores / denominator
     scores = scores * routed_scaling_factor
 
@@ -437,8 +436,6 @@ class DeepseekV3Model(nn.Module):
 
         pipeline_rank = self.pipeline_rank
         pipeline_size = self.pipeline_size
-        # Hack to avoid time-outs during prompt-processing
-        dist_stream = mx.cpu if h.shape[1] > 1 else mx.gpu
         if mask is None:
             mask = create_attention_mask(h, cache)
 
@@ -448,19 +445,17 @@ class DeepseekV3Model(nn.Module):
         # Receive from the previous process in the pipeline
 
         if pipeline_rank < pipeline_size - 1:
-            h = mx.distributed.recv_like(h, (pipeline_rank + 1), stream=dist_stream)
+            h = mx.distributed.recv_like(h, (pipeline_rank + 1))
 
         for i in range(self.num_layers):
             h = self.layers[self.start_idx + i](h, mask, cache[i])
 
         # Send to the next process in the pipeline
         if pipeline_rank != 0:
-            h = mx.distributed.send(
-                h, (pipeline_rank - 1) % pipeline_size, stream=dist_stream
-            )
+            h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
 
         # Broadcast h while keeping it in the graph
-        h = mx.distributed.all_gather(h, stream=dist_stream)[: h.shape[0]]
+        h = mx.distributed.all_gather(h)[: h.shape[0]]
 
         return self.norm(h)
 
@@ -484,6 +479,7 @@ class Model(nn.Module):
 
     def sanitize(self, weights):
         def dequant(weight, scale_inv):
+            dtype = weight.dtype
             bs = 128  # block size
             m, n = weight.shape
             pad_bottom = (-m) % bs
@@ -492,11 +488,10 @@ class Model(nn.Module):
             weight = weight.reshape(
                 ((m + pad_bottom) // bs, bs, (n + pad_side) // bs, bs)
             )
-            scale_inv = scale_inv.astype(weight.dtype)
             weight = (weight * scale_inv[:, None, :, None]).reshape(
                 m + pad_bottom, n + pad_side
             )
-            return weight[:m, :n]
+            return weight[:m, :n].astype(dtype)
 
         # Dequantize
         new_weights = {}
@@ -533,3 +528,9 @@ class Model(nn.Module):
     @property
     def layers(self):
         return self.model.layers[self.model.start_idx : self.model.end_idx]
+
+    def cast_predicate(self):
+        def predicate(k):
+            return "e_score_correction_bias" not in k
+
+        return predicate

mlx_lm/models/gemma3_text.py

@@ -1,12 +1,13 @@
 # Copyright © 2025 Apple Inc.
 
 from dataclasses import dataclass
+from functools import partial
 from typing import Any, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .cache import KVCache, RotatingKVCache
 
 
@@ -88,9 +89,8 @@ class Attention(nn.Module):
         # Sliding window
         if isinstance(mask, mx.array) and mask.shape[-1] != keys.shape[-2]:
             mask = mask[..., -keys.shape[-2] :]
-
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        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)
@@ -117,6 +117,16 @@ class MLP(nn.Module):
         return self.down_proj(nn.gelu_approx(self.gate_proj(x)) * self.up_proj(x))
 
 
+@partial(mx.compile, shapeless=True)
+def clip_residual(x, y):
+    if x.dtype != mx.float16:
+        return x + y
+    bound = mx.finfo(mx.float16).max
+    return mx.clip(x.astype(mx.float32) + y.astype(mx.float32), -bound, bound).astype(
+        mx.float16
+    )
+
+
 class TransformerBlock(nn.Module):
     def __init__(self, args: ModelArgs, layer_idx: int):
         super().__init__()
@@ -140,9 +150,9 @@ class TransformerBlock(nn.Module):
         cache: Optional[Any] = None,
     ) -> mx.array:
         r = self.self_attn(self.input_layernorm(x), mask, cache)
-        h = x + self.post_attention_layernorm(r)
+        h = clip_residual(x, self.post_attention_layernorm(r))
         r = self.mlp(self.pre_feedforward_layernorm(h))
-        out = h + self.post_feedforward_layernorm(r)
+        out = clip_residual(h, self.post_feedforward_layernorm(r))
         return out
 
 

mlx_lm/models/glm4.py

@@ -0,0 +1,183 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    attention_bias: bool
+    head_dim: int
+    rms_norm_eps: float
+    vocab_size: int
+    num_key_value_heads: int
+    partial_rotary_factor: float
+    rope_theta: float
+    rope_traditional: bool = True
+    max_position_embeddings: int = 32768
+
+
+class Glm4MLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.gate_up_proj = nn.Linear(
+            args.hidden_size, 2 * args.intermediate_size, bias=False
+        )
+        self.down_proj = nn.Linear(args.intermediate_size, args.hidden_size, bias=False)
+
+    def __call__(self, x) -> mx.array:
+        x = self.gate_up_proj(x)
+        gate, up_states = mx.split(x, 2, axis=-1)
+        return self.down_proj(nn.silu(gate) * up_states)
+
+
+class Glm4Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.head_dim = getattr(
+            args, "head_dim", args.hidden_size // args.num_attention_heads
+        )
+        self.n_heads = args.num_attention_heads
+        self.n_kv_heads = args.num_key_value_heads
+        self.scale = self.head_dim**-0.5
+
+        self.q_proj = nn.Linear(
+            args.hidden_size,
+            args.num_attention_heads * self.head_dim,
+            bias=args.attention_bias,
+        )
+        self.k_proj = nn.Linear(
+            args.hidden_size,
+            args.num_key_value_heads * self.head_dim,
+            bias=args.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            args.hidden_size,
+            args.num_key_value_heads * self.head_dim,
+            bias=args.attention_bias,
+        )
+        self.o_proj = nn.Linear(
+            args.num_attention_heads * self.head_dim, args.hidden_size, bias=False
+        )
+
+        self.rope = nn.RoPE(
+            dims=int(self.head_dim * args.partial_rotary_factor),
+            base=args.rope_theta,
+            traditional=args.rope_traditional,
+        )
+
+    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).transpose(0, 2, 1, 3)
+        keys = 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, scale=self.scale, mask=mask
+        )
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class Glm4DecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.self_attn = Glm4Attention(args=args)
+
+        self.mlp = Glm4MLP(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
+        )
+        self.post_self_attn_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+        self.post_mlp_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:
+        x = x + self.post_self_attn_layernorm(
+            self.self_attn(self.input_layernorm(x), mask, cache)
+        )
+        residual = x
+        x = (
+            self.post_mlp_layernorm(self.mlp(self.post_attention_layernorm(x)))
+            + residual
+        )
+        return x
+
+
+class Glm4Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            Glm4DecoderLayer(args=args) for _ 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: Optional[mx.array] = None,
+        cache: Optional[Any] = 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, cache=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 = Glm4Model(args)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ):
+        out = self.model(inputs, mask, cache)
+        return self.lm_head(out)
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/gpt2.py

@@ -1,7 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.
 
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -133,14 +133,15 @@ class GPT2Model(nn.Module):
 
         hidden_states = self.wte(inputs)
 
-        mask = None
-        if hidden_states.shape[1] > 1:
+        offset = 0
+        if cache is not None and len(cache) > 0 and cache[0] is not None:
+            offset = cache[0].offset
 
-            position_ids = mx.array(np.arange(L))
-            hidden_states += self.wpe(position_ids)
+        position_ids = mx.arange(offset, offset + L)
+        hidden_states += self.wpe(position_ids)
 
-            if mask is None:
-                mask = create_attention_mask(hidden_states, cache)
+        if mask is None:
+            mask = create_attention_mask(hidden_states, cache)
 
         if cache is None:
             cache = [None] * len(self.h)

mlx_lm/models/kimi_vl.py

@@ -0,0 +1,120 @@
+# Copyright © 2024 Apple Inc.
+
+import math
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, Dict, Optional, Tuple, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .deepseek_v3 import DeepseekV3Model
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class TextArgs(BaseModelArgs):
+    vocab_size: int = 102400
+    hidden_size: int = 4096
+    intermediate_size: int = 11008
+    moe_intermediate_size: int = 1407
+    num_hidden_layers: int = 30
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 32
+    n_shared_experts: Optional[int] = None
+    n_routed_experts: Optional[int] = None
+    routed_scaling_factor: float = 1.0
+    kv_lora_rank: int = 512
+    q_lora_rank: int = 1536
+    qk_rope_head_dim: int = 64
+    v_head_dim: int = 128
+    qk_nope_head_dim: int = 128
+    topk_method: str = "noaux_tc"
+    scoring_func: str = "sigmoid"
+    norm_topk_prob: bool = True
+    n_group: Optional[int] = None
+    topk_group: Optional[int] = None
+    num_experts_per_tok: Optional[int] = None
+    moe_layer_freq: int = 1
+    first_k_dense_replace: int = 0
+    max_position_embeddings: int = 2048
+    rms_norm_eps: float = 1e-6
+    rope_theta: float = 10000.0
+    rope_scaling: Dict = None
+    attention_bias: bool = False
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    text_config: Union[TextArgs, dict]
+    model_type: str
+
+    def __post_init__(self):
+        self.text_config = TextArgs.from_dict(self.text_config)
+
+
+class LanguageModel(nn.Module):
+    def __init__(self, config: TextArgs):
+        super().__init__()
+        self.args = config
+        self.model = DeepseekV3Model(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)
+
+
+class Model(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.args = config
+        self.model_type = config.model_type
+        self.language_model = LanguageModel(config.text_config)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+        mask: Optional[mx.array] = None,
+    ):
+        return self.language_model(inputs, cache, mask)
+
+    def sanitize(self, weights):
+        def keep(key):
+            return (
+                "vision_tower" not in key
+                and "rotary_emb" not in key
+                and "multi_modal_projector" not in key
+            )
+
+        weights = {k: v for k, v in weights.items() if keep(k)}
+        # Stack experts
+        for l in range(self.args.text_config.num_hidden_layers):
+            prefix = f"language_model.model.layers.{l}"
+            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 = [
+                            weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}")
+                            for e in range(self.args.text_config.n_routed_experts)
+                        ]
+                        weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
+
+        return weights
+
+    @property
+    def layers(self):
+        return self.language_model.model.layers
+
+    def cast_predicate(self):
+        def predicate(k):
+            return "e_score_correction_bias" not in k
+
+        return predicate

mlx_lm/models/llama4.py

@@ -0,0 +1,333 @@
+# 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 ChunkedKVCache, KVCache
+from .rope_utils import initialize_rope
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class TextArgs(BaseModelArgs):
+    attention_bias: bool
+    attention_chunk_size: int
+    head_dim: int
+    hidden_act: str
+    hidden_size: int
+    interleave_moe_layer_step: int
+    intermediate_size: int
+    intermediate_size_mlp: int
+    max_position_embeddings: int
+    model_type: str
+    num_attention_heads: int
+    num_experts_per_tok: int
+    num_hidden_layers: int
+    num_key_value_heads: int
+    num_local_experts: int
+    rms_norm_eps: float
+    rope_scaling: Any
+    rope_theta: float
+    use_qk_norm: bool
+    vocab_size: int
+    attn_temperature_tuning: int = 4
+    floor_scale: int = 8192
+    attn_scale: float = 0.1
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    text_config: Union[TextArgs, dict]
+    model_type: str
+
+    def __post_init__(self):
+        self.text_config = TextArgs.from_dict(self.text_config)
+
+
+class Attention(nn.Module):
+    def __init__(self, args: TextArgs, layer_idx: int):
+        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.use_rope = int((layer_idx + 1) % 4 != 0)  # rope unused for dense layers
+        self.attn_temperature_tuning = args.attn_temperature_tuning
+        self.floor_scale = args.floor_scale
+        self.attn_scale = args.attn_scale
+
+        self.head_dim = head_dim = args.head_dim or args.hidden_size // n_heads
+
+        self.scale = head_dim**-0.5
+        if hasattr(args, "attention_bias"):
+            attention_bias = args.attention_bias
+        else:
+            attention_bias = False
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=attention_bias)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=attention_bias)
+
+        self.use_qk_norm = args.use_qk_norm and self.use_rope
+
+        if self.use_rope:
+            self.rope = initialize_rope(
+                head_dim,
+                args.rope_theta,
+                traditional=True,
+                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 = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        keys = 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:
+            offset = cache.offset
+        else:
+            offset = 0
+
+        if self.use_rope:
+            queries = self.rope(queries, offset=offset)
+            keys = self.rope(keys, offset=offset)
+
+        if self.use_qk_norm:
+            queries = mx.fast.rms_norm(queries, weight=None, eps=1e-6)
+            keys = mx.fast.rms_norm(keys, weight=None, eps=1e-6)
+
+        if self.attn_temperature_tuning and not self.use_rope:
+            attn_scales = (
+                mx.log(
+                    mx.floor(mx.arange(offset + 1, offset + L + 1) / self.floor_scale)
+                    + 1.0
+                )
+                * self.attn_scale
+                + 1.0
+            )
+            attn_scales = attn_scales[:, None]
+            queries = (queries * attn_scales).astype(queries.dtype)
+
+        if cache is not None:
+            keys, values = cache.update_and_fetch(keys, values)
+
+        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, args: ModelArgs, intermediate_size: int = None):
+        super().__init__()
+
+        dim = args.hidden_size
+        hidden_dim = intermediate_size or args.intermediate_size
+
+        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 MoE(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.top_k = args.num_experts_per_tok
+        self.num_experts = args.num_local_experts
+        self.experts = SwitchGLU(
+            args.hidden_size, args.intermediate_size, self.num_experts
+        )
+        self.router = nn.Linear(args.hidden_size, args.num_local_experts, bias=False)
+        self.shared_expert = MLP(args)
+
+    def __call__(self, x) -> mx.array:
+        logits = self.router(x)
+        k = self.top_k
+        indices = mx.argpartition(-logits, kth=k - 1, axis=-1)[..., :k]
+        scores = mx.take_along_axis(logits, indices, axis=-1)
+        scores = mx.sigmoid(scores.astype(mx.float32)).astype(x.dtype)
+
+        out = self.experts(x * scores, indices).squeeze(2)
+        return out + self.shared_expert(x)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: TextArgs, layer_idx: int):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.hidden_size = args.hidden_size
+        self.self_attn = Attention(args, layer_idx)
+        self.is_moe_layer = (layer_idx % args.interleave_moe_layer_step) == (
+            args.interleave_moe_layer_step - 1
+        )
+        if self.is_moe_layer:
+            self.feed_forward = MoE(args)
+        else:
+            self.feed_forward = MLP(args, args.intermediate_size_mlp)
+
+        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
+        )
+        self.args = args
+
+    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.feed_forward(self.post_attention_layernorm(h))
+        out = h + r
+        return out
+
+
+class LlamaModel(nn.Module):
+    def __init__(self, args: TextArgs):
+        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)
+        self.layers = [TransformerBlock(args, i) for i in range(args.num_hidden_layers)]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.attention_chunk_size = args.attention_chunk_size
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        h = self.embed_tokens(inputs)
+
+        if cache is not None:
+            for idx, c in enumerate(cache):
+                if (idx + 1) % 4 != 0:
+                    c.maybe_trim_front()
+            start = cache[0].start_position
+            offset = cache[0].offset
+        else:
+            start = 0
+            offset = 0
+        end = offset + h.shape[1]
+        linds = mx.arange(start, end)
+        rinds = mx.arange(offset, end)[:, None]
+        block_pos = mx.abs(
+            (linds // self.attention_chunk_size) - (rinds // self.attention_chunk_size)
+        )
+        token_pos = linds <= rinds
+        chunk_mask = (block_pos == 0) & token_pos
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+        else:
+            chunk_mask &= mask
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for idx, (layer, c) in enumerate(zip(self.layers, cache)):
+            use_chunked_attention = (idx + 1) % 4 != 0
+            if use_chunked_attention:
+                local_mask = chunk_mask
+            else:
+                local_mask = mask
+            h = layer(h, local_mask, cache=c)
+
+        return self.norm(h)
+
+
+class LanguageModel(nn.Module):
+    def __init__(self, args: TextArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = LlamaModel(self.args)
+        self.lm_head = nn.Linear(
+            self.args.hidden_size, self.args.vocab_size, bias=False
+        )
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        out = self.model(inputs, mask, cache)
+        return self.lm_head(out)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.language_model = LanguageModel(args.text_config)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        return self.language_model(inputs, mask, cache)
+
+    def sanitize(self, weights):
+        def to_remove(k):
+            return "vision_model" in k or "multi_modal_projector" in k
+
+        # Remove vision weights
+        weights = {k: v for k, v in weights.items() if not to_remove(k)}
+
+        # Rename expert weights for SwitchGLU
+        for l in range(self.args.text_config.num_hidden_layers):
+            prefix = f"language_model.model.layers.{l}.feed_forward.experts"
+            if f"{prefix}.gate_up_proj" in weights:
+                v = weights.pop(f"{prefix}.gate_up_proj")
+                gate_k = f"{prefix}.gate_proj.weight"
+                up_k = f"{prefix}.up_proj.weight"
+                gate_proj, up_proj = mx.split(v, 2, axis=-1)
+                weights[gate_k] = mx.swapaxes(gate_proj, 1, 2)
+                weights[up_k] = mx.swapaxes(up_proj, 1, 2)
+            if f"{prefix}.down_proj" in weights:
+                down_proj = weights.pop(f"{prefix}.down_proj")
+                weights[f"{prefix}.down_proj.weight"] = mx.swapaxes(down_proj, 1, 2)
+        return weights
+
+    @property
+    def layers(self):
+        return self.language_model.model.layers
+
+    def make_cache(self):
+        chunk_size = self.args.text_config.attention_chunk_size
+        caches = []
+        for i in range(len(self.layers)):
+            if (i + 1) % 4 != 0:
+                caches.append(ChunkedKVCache(chunk_size))
+            else:
+                caches.append(KVCache())
+        return caches

mlx_lm/models/mimo.py

@@ -0,0 +1,194 @@
+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 .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 = 32768
+    rope_theta: float = 10000.0
+    rope_traditional: bool = False
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = False
+    num_nextn_predict_layers: int = 2
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        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.hidden_size // n_heads
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=True)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=True)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=True)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
+
+        self.rope = initialize_rope(
+            head_dim,
+            base=args.rope_theta,
+            traditional=args.rope_traditional,
+            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 = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        keys = 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, 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):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.hidden_size = args.hidden_size
+        self.self_attn = Attention(args)
+        self.mlp = MLP(args.hidden_size, args.intermediate_size)
+        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
+        )
+        self.args = args
+
+    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))
+        out = h + r
+        return out
+
+
+class MiMoModel(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.num_nextn_predict_layers = args.num_nextn_predict_layers
+
+        assert self.vocab_size > 0
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            TransformerBlock(args=args) for _ 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)
+
+        h = self.norm(h)
+
+        return h
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = MiMoModel(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 sanitize(self, weights):
+        if self.args.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
+
+        return {
+            k: v
+            for k, v in weights.items()
+            if "self_attn.rotary_emb.inv_freq" not in k
+            and not k.startswith("model.mtp_layers.")
+        }
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/nemotron-nas.py

@@ -0,0 +1,386 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass, field
+from dataclasses import fields as dataclass_fields
+from typing import Any, Dict, List, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
+
+
+@dataclass(frozen=True)
+class AttentionConfig:
+    no_op: bool = False
+    replace_with_linear: bool = False
+    sparsify: Optional[list[str]] = None
+    n_heads_in_group: Optional[int] = None  # GQA group size
+    window_length: Optional[int] = None  # Not directly used here, placeholder
+    num_sink_tokens: Optional[int] = None  # Not directly used here, placeholder
+    use_prefill_window_in_sink_attention: bool = (
+        False  # Not directly used here, placeholder
+    )
+    unshifted_sink: bool = False  # Not directly used here, placeholder
+
+    def __post_init__(self):
+        # Ensure consistency: If no-op or linear, other attn params are irrelevant
+        if self.no_op or self.replace_with_linear:
+            # Use object.__setattr__ because the dataclass is frozen
+            object.__setattr__(self, "n_heads_in_group", None)
+            object.__setattr__(self, "window_length", None)
+            object.__setattr__(self, "num_sink_tokens", None)
+        # If it's a standard attention block, n_heads_in_group must be provided
+        elif not self.no_op:
+            if self.n_heads_in_group is None:
+                raise ValueError(
+                    "n_heads_in_group must be specified for active attention blocks"
+                )
+            if self.n_heads_in_group <= 0:
+                raise ValueError(
+                    f"n_heads_in_group must be positive, got {self.n_heads_in_group}"
+                )
+
+
+@dataclass(frozen=True)
+class FFNConfig:
+    no_op: bool = False
+    replace_with_linear: bool = False
+    sparsify: Optional[list[str]] = None
+    ffn_mult: Optional[float] = None
+
+    def __post_init__(self):
+        # Ensure consistency: If no-op or linear, ffn_mult is irrelevant
+        if self.no_op or self.replace_with_linear:
+            object.__setattr__(self, "ffn_mult", None)
+        # If it's a standard FFN block, ffn_mult must be provided
+        elif not self.no_op:
+            if self.ffn_mult is None:
+                raise ValueError("ffn_mult must be specified for active FFN blocks")
+            # Round to prevent potential floating point inconsistencies if needed
+            object.__setattr__(self, "ffn_mult", round(self.ffn_mult, 6))
+
+
+@dataclass(frozen=True)
+class BlockConfig:
+    attention: AttentionConfig
+    ffn: FFNConfig
+
+    @classmethod
+    def from_dict(cls, data: dict):
+        # Helper to create BlockConfig from a dictionary (e.g., loaded from JSON)
+        attn_conf = AttentionConfig(**data.get("attention", {}))
+        ffn_conf = FFNConfig(**data.get("ffn", {}))
+        return cls(attention=attn_conf, ffn=ffn_conf)
+
+
+def _find_multiple(n: int, k: int) -> int:
+    """Finds the smallest multiple of k greater than or equal to n."""
+    if n % k == 0:
+        return n
+    return n + k - (n % k)
+
+
+def _ffn_mult_to_intermediate_size(ffn_mult: float, n_embd: int) -> int:
+    """Calculates intermediate size based on multiplier, rounding up to multiple of 256."""
+    intermediate_size = int(2 * ffn_mult * n_embd / 3)
+    return _find_multiple(intermediate_size, 256)
+
+
+# Activation function mapping
+_ACT2FN = {
+    "silu": nn.silu,
+    "relu": nn.relu,
+    "gelu": nn.gelu,
+    "gelu_new": nn.gelu_approx,
+    "gelu_fast": nn.gelu_approx,
+}
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str = "nemotron-nas"
+    hidden_size: int = 8192
+    num_hidden_layers: int = 80
+    num_attention_heads: int = 64
+    rms_norm_eps: float = 1e-5
+    vocab_size: int = 128256
+    block_configs: list = field(default_factory=list)  # List of BlockConfig or dicts
+    hidden_act: str = "silu"
+    attention_bias: bool = False
+    mlp_bias: bool = False
+    rope_theta: float = 500000.0
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    max_position_embeddings: int = 131072
+    tie_word_embeddings: bool = False
+
+    def __post_init__(self):
+        # Automatically parse block_configs if they are loaded as dicts
+        if self.block_configs and isinstance(self.block_configs[0], dict):
+            self.block_configs = [
+                BlockConfig.from_dict(conf) for conf in self.block_configs
+            ]
+
+        if len(self.block_configs) != self.num_hidden_layers:
+            raise ValueError(
+                f"Number of block_configs ({len(self.block_configs)}) must match "
+                f"num_hidden_layers ({self.num_hidden_layers})"
+            )
+
+        # Basic validation for RoPE scaling if provided
+        if self.rope_scaling:
+            if "factor" not in self.rope_scaling:
+                raise ValueError("rope_scaling must contain 'factor'")
+            rope_type = self.rope_scaling.get("rope_type")
+            if rope_type is None:
+                raise ValueError("rope_scaling must contain 'rope_type'")
+
+        # Validate individual block configs (post_init in dataclasses already does some)
+        for i, block_conf in enumerate(self.block_configs):
+            attn_conf = block_conf.attention
+            if not attn_conf.no_op and not attn_conf.replace_with_linear:
+                if self.num_attention_heads % attn_conf.n_heads_in_group != 0:
+                    raise ValueError(
+                        f"Layer {i}: num_attention_heads ({self.num_attention_heads}) "
+                        f"must be divisible by n_heads_in_group ({attn_conf.n_heads_in_group})"
+                    )
+
+
+class Attention(nn.Module):
+    """Standard GQA Attention mechanism for layers that use it."""
+
+    def __init__(self, args: ModelArgs, attention_config: AttentionConfig):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = n_heads // attention_config.n_heads_in_group
+
+        self.head_dim = head_dim = args.hidden_size // n_heads
+        if (self.head_dim * n_heads) != dim:
+            raise ValueError(
+                f"hidden_size ({dim}) must be divisible by num_attention_heads ({n_heads})"
+            )
+
+        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=args.attention_bias)
+
+        # Initialize RoPE based on global config
+        self.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            False,  # Llama uses traditional=False
+            args.rope_scaling,
+            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 = queries.reshape(B, L, self.n_heads, self.head_dim).transpose(
+            0, 2, 1, 3
+        )
+        keys = keys.reshape(B, L, self.n_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, self.head_dim).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):
+    """Standard Feed-Forward Network for layers that use it."""
+
+    def __init__(self, args: ModelArgs, ffn_config: FFNConfig):
+        super().__init__()
+
+        dim = args.hidden_size
+        # Calculate intermediate dim based on layer's specific config
+        hidden_dim = _ffn_mult_to_intermediate_size(ffn_config.ffn_mult, dim)
+
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=args.mlp_bias)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=args.mlp_bias)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=args.mlp_bias)
+
+        try:
+            self.act_fn = _ACT2FN[args.hidden_act]
+        except KeyError:
+            raise ValueError(f"Unknown activation function: {args.hidden_act}")
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+
+
+class LinearSubblockReplacement(nn.Module):
+    """A simple linear layer used to replace Attention or MLP blocks."""
+
+    def __init__(self, hidden_size: int, bias: bool):
+        super().__init__()
+        self.linear = nn.Linear(hidden_size, hidden_size, bias=bias)
+
+    def __call__(self, x: mx.array, *args, **kwargs) -> mx.array:
+        # Accepts potential extra args (like mask, cache) but ignores them
+        return self.linear(x)
+
+
+class TransformerBlock(nn.Module):
+    """A single transformer block, potentially heterogeneous based on config."""
+
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.hidden_size = args.hidden_size
+        # Get the specific configuration for this layer
+        block_config = args.block_configs[layer_idx]
+        self.attention_config = block_config.attention
+        self.ffn_config = block_config.ffn
+
+        # Conditionally initialize Input LayerNorm (needed unless Attention is no-op)
+        if not self.attention_config.no_op:
+            self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        else:
+            self.input_layernorm = None
+
+        # Conditionally initialize Attention block
+        if self.attention_config.no_op:
+            self.self_attn = None
+        elif self.attention_config.replace_with_linear:
+            self.self_attn = LinearSubblockReplacement(
+                args.hidden_size, args.attention_bias
+            )
+        else:
+            # Standard attention for this layer
+            self.self_attn = Attention(args, self.attention_config)
+
+        # Conditionally initialize Post-Attention LayerNorm (needed unless FFN is no-op)
+        if not self.ffn_config.no_op:
+            self.post_attention_layernorm = nn.RMSNorm(
+                args.hidden_size, eps=args.rms_norm_eps
+            )
+        else:
+            self.post_attention_layernorm = None
+
+        # Conditionally initialize MLP block
+        if self.ffn_config.no_op:
+            self.mlp = None
+        elif self.ffn_config.replace_with_linear:
+            self.mlp = LinearSubblockReplacement(args.hidden_size, args.mlp_bias)
+        else:
+            # Standard MLP for this layer
+            self.mlp = MLP(args, self.ffn_config)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+
+        # Attention part (Input Norm -> Attention -> Residual)
+        if self.self_attn is not None:
+            residual = x
+            h = self.input_layernorm(x)
+            attn_out = self.self_attn(h, mask=mask, cache=cache)
+            x = residual + attn_out
+
+        # MLP part (Post-Attention Norm -> MLP -> Residual)
+        if self.mlp is not None:
+            residual = x
+            h = self.post_attention_layernorm(x)
+            mlp_out = self.mlp(h)
+            x = residual + mlp_out
+
+        return x
+
+
+class NemotronNASModel(nn.Module):
+    """The core Nemotron-NAS style transformer model."""
+
+    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 = [
+            TransformerBlock(args=args, layer_idx=i)
+            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: Optional[mx.array] = None,
+        cache: Optional[List[Any]] = 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 i, layer in enumerate(self.layers):
+            h = layer(h, mask, cache=cache[i])
+
+        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 = NemotronNASModel(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+        else:
+            self.lm_head = None
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask=None,
+        cache=None,
+    ):
+        out = self.model(inputs, mask=mask, cache=cache)
+        if self.args.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
+        return out
+
+    def sanitize(self, weights):
+        if self.args.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
+        return weights
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/plamo2.py

@@ -53,16 +53,6 @@ class RMSNorm(nn.Module):
         )
 
 
-def _rms_norm(hidden_states: mx.array, eps: float) -> mx.array:
-    input_dtype = hidden_states.dtype
-    hidden_states = hidden_states.astype(mx.float32)
-    variance = mx.power(hidden_states, 2).mean(-1, keepdims=True)
-    hidden_states = hidden_states * mx.rsqrt(variance + eps)
-    hidden_states = hidden_states.astype(input_dtype)
-
-    return hidden_states
-
-
 def get_initial_dt_bias(num_heads: int) -> mx.array:
     dt_min = 0.001
     dt_max = 0.1
@@ -401,8 +391,8 @@ class Attention(nn.Module):
         k = k.reshape(B, T, self.k_num_heads, self.qk_dim).transpose(0, 2, 1, 3)
         v = v.reshape(B, T, self.v_num_heads, self.v_dim).transpose(0, 2, 1, 3)
 
-        q = _rms_norm(q, 1e-6) * self.q_weight[:, None]
-        k = _rms_norm(k, 1e-6) * self.k_weight[:, None]
+        q = mx.fast.rms_norm(q, weight=None, eps=1e-6) * self.q_weight[:, None]
+        k = mx.fast.rms_norm(k, weight=None, eps=1e-6) * self.k_weight[:, None]
 
         if cache is not None:
             q = self.rope(q, offset=cache.offset)

mlx_lm/models/qwen3.py

@@ -0,0 +1,189 @@
+# 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 .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: Optional[Dict[str, Union[float, str]]] = None
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        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.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:
+            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, 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):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.hidden_size = args.hidden_size
+        self.self_attn = Attention(args)
+        self.mlp = MLP(args.hidden_size, args.intermediate_size)
+        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
+        )
+        self.args = args
+
+    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))
+        out = h + r
+        return out
+
+
+class Qwen3Model(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)
+        self.layers = [
+            TransformerBlock(args=args) for _ 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 = Qwen3Model(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 sanitize(self, weights):
+        if self.args.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
+        return weights
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/qwen3_moe.py

@@ -0,0 +1,243 @@
+# Copyright © 2025 Apple Inc.
+
+import math
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    num_experts_per_tok: int
+    num_experts: int
+    num_experts_per_tok: int
+    decoder_sparse_step: int
+    mlp_only_layers: List[int]
+    moe_intermediate_size: int
+    rms_norm_eps: float
+    vocab_size: int
+    num_key_value_heads: int
+    head_dim: int
+    rope_theta: float
+    tie_word_embeddings: bool
+    max_position_embeddings: int
+    norm_topk_prob: bool
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        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 = getattr(
+            args, "head_dim", args.hidden_size // args.num_attention_heads
+        )
+        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.rope = nn.RoPE(
+            head_dim,
+            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)
+
+        # Prepare the queries, keys and values for the attention computation
+        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:
+            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, 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 Qwen3MoeSparseMoeBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        dim = args.hidden_size
+        intermediate_size = args.moe_intermediate_size
+
+        self.num_experts = num_experts = args.num_experts
+        self.top_k = args.num_experts_per_tok
+        self.norm_topk_prob = args.norm_topk_prob
+
+        self.gate = nn.Linear(dim, num_experts, bias=False)
+        self.switch_mlp = SwitchGLU(dim, intermediate_size, num_experts)
+
+    def __call__(
+        self,
+        x: mx.array,
+    ):
+        gates = self.gate(x)
+        gates = mx.softmax(gates, axis=-1, precise=True)
+
+        k = self.top_k
+        inds = mx.stop_gradient(mx.argpartition(-gates, kth=k - 1, axis=-1)[..., :k])
+        scores = mx.take_along_axis(gates, inds, axis=-1)
+        if self.norm_topk_prob:
+            scores /= mx.sum(scores, axis=-1, keepdims=True)
+
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2)
+
+        return y
+
+
+class Qwen3MoeDecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.hidden_size = args.hidden_size
+        self.self_attn = Attention(args, layer_idx)
+
+        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
+        )
+        self.args = args
+
+        if (layer_idx not in args.mlp_only_layers) and (
+            args.num_experts > 0 and (layer_idx + 1) % args.decoder_sparse_step == 0
+        ):
+            self.mlp = Qwen3MoeSparseMoeBlock(args)
+        else:
+            self.mlp = MLP(args.hidden_size, args.intermediate_size)
+
+    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))
+        out = h + r
+        return out
+
+
+class Qwen3MoeModel(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)
+        self.layers = [
+            Qwen3MoeDecoderLayer(args=args, layer_idx=i)
+            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 = Qwen3MoeModel(args)
+        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)
+        return self.lm_head(out)
+
+    def sanitize(self, weights):
+        if "model.layers.0.mlp.experts.0.up_proj.weight" not in weights:
+            return weights
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}"
+            for n in ["up_proj", "down_proj", "gate_proj"]:
+                if f"{prefix}.mlp.experts.0.{n}.weight" in weights:
+                    to_join = [
+                        weights.pop(f"{prefix}.mlp.experts.{e}.{n}.weight")
+                        for e in range(self.args.num_experts)
+                    ]
+                    weights[f"{prefix}.mlp.switch_mlp.{n}.weight"] = mx.stack(to_join)
+        return weights
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/switch_layers.py

@@ -6,6 +6,21 @@ import mlx.core as mx
 import mlx.nn as nn
 
 
+def _gather_sort(x, indices):
+    *_, M = indices.shape
+    indices = indices.flatten()
+    order = mx.argsort(indices)
+    inv_order = mx.argsort(order)
+    return x.flatten(0, -3)[order // M], indices[order], inv_order
+
+
+def _scatter_unsort(x, inv_order, shape=None):
+    x = x[inv_order]
+    if shape is not None:
+        x = mx.unflatten(x, 0, shape)
+    return x
+
+
 class QuantizedSwitchLinear(nn.Module):
     def __init__(
         self,
@@ -56,7 +71,7 @@ class QuantizedSwitchLinear(nn.Module):
     def num_experts(self):
         return self.weight.shape[0]
 
-    def __call__(self, x, indices):
+    def __call__(self, x, indices, sorted_indices=False):
         x = mx.gather_qmm(
             x,
             self["weight"],
@@ -66,6 +81,7 @@ class QuantizedSwitchLinear(nn.Module):
             transpose=True,
             group_size=self.group_size,
             bits=self.bits,
+            sorted_indices=sorted_indices,
         )
         if "bias" in self:
             x = x + mx.expand_dims(self["bias"][indices], -2)
@@ -99,8 +115,13 @@ class SwitchLinear(nn.Module):
     def num_experts(self):
         return self.weight.shape[0]
 
-    def __call__(self, x, indices):
-        x = mx.gather_mm(x, self["weight"].swapaxes(-1, -2), rhs_indices=indices)
+    def __call__(self, x, indices, sorted_indices=False):
+        x = mx.gather_mm(
+            x,
+            self["weight"].swapaxes(-1, -2),
+            rhs_indices=indices,
+            sorted_indices=sorted_indices,
+        )
         if "bias" in self:
             x = x + mx.expand_dims(self["bias"][indices], -2)
         return x
@@ -122,7 +143,7 @@ class SwitchGLU(nn.Module):
         input_dims: int,
         hidden_dims: int,
         num_experts: int,
-        activation=nn.silu,
+        activation=nn.SiLU(),
         bias: bool = False,
     ):
         super().__init__()
@@ -135,9 +156,24 @@ class SwitchGLU(nn.Module):
     def __call__(self, x, indices) -> mx.array:
         x = mx.expand_dims(x, (-2, -3))
 
-        x_up = self.up_proj(x, indices)
-        x_gate = self.gate_proj(x, indices)
-        x = self.down_proj(self.activation(x_gate) * x_up, indices)
+        # When we have many tokens, then sort them to make sure that the access
+        # of different experts is in order.
+        do_sort = indices.size >= 64
+        idx = indices
+        inv_order = None
+        if do_sort:
+            x, idx, inv_order = _gather_sort(x, indices)
+
+        x_up = self.up_proj(x, idx, sorted_indices=do_sort)
+        x_gate = self.gate_proj(x, idx, sorted_indices=do_sort)
+        x = self.down_proj(
+            self.activation(x_gate) * x_up,
+            idx,
+            sorted_indices=do_sort,
+        )
+
+        if do_sort:
+            x = _scatter_unsort(x, inv_order, indices.shape)
 
         return x.squeeze(-2)
 
@@ -148,7 +184,7 @@ class SwitchMLP(nn.Module):
         input_dims: int,
         hidden_dims: int,
         num_experts: int,
-        activation=nn.gelu_approx,
+        activation=nn.GELU(approx="precise"),
         bias: bool = False,
     ):
         super().__init__()
@@ -160,8 +196,19 @@ class SwitchMLP(nn.Module):
     def __call__(self, x, indices) -> mx.array:
         x = mx.expand_dims(x, (-2, -3))
 
-        x = self.fc1(x, indices)
+        # When we have many tokens, then sort them to make sure that the access
+        # of different experts is in order.
+        do_sort = indices.size >= 64
+        idx = indices
+        inv_order = None
+        if do_sort:
+            x, idx, inv_order = _gather_sort(x, indices)
+
+        x = self.fc1(x, idx, sorted_indices=do_sort)
         x = self.activation(x)
-        x = self.fc2(x, indices)
+        x = self.fc2(x, idx, sorted_indices=do_sort)
+
+        if do_sort:
+            x = _scatter_unsort(x, inv_order, indices.shape)
 
         return x.squeeze(-2)

mlx_lm/qat.py

@@ -0,0 +1,339 @@
+# Copyright © 2025 Apple Inc.
+
+import argparse
+import copy
+import glob
+import shutil
+import time
+import types
+from pathlib import Path
+
+import mlx.core as mx
+import mlx.nn as nn
+import mlx.optimizers as optimizers
+import numpy as np
+from mlx.utils import tree_flatten, tree_map, tree_map_with_path
+
+from mlx_lm.tokenizer_utils import TokenizerWrapper
+from mlx_lm.tuner.datasets import load_dataset
+from mlx_lm.tuner.trainer import iterate_batches
+from mlx_lm.tuner.utils import print_trainable_parameters
+from mlx_lm.utils import (
+    create_model_card,
+    fetch_from_hub,
+    get_model_path,
+    quantize_model,
+    save_config,
+    save_weights,
+)
+
+
+class StraightThroughQuantizedEmbedding(nn.Module):
+    def __init__(
+        self,
+        num_embeddings: int,
+        dims: int,
+        group_size: int = 64,
+        bits: int = 4,
+    ):
+        super().__init__()
+
+        # Quantization config
+        self.group_size = group_size
+        self.bits = bits
+
+        # Initialize the quantized weight
+        self.weight = mx.zeros(shape=(num_embeddings, dims))
+        self.num_embeddings = num_embeddings
+        self.dims = dims
+
+    def __call__(self, x):
+        w, s, b = mx.quantize(self.weight, self.group_size, self.bits)
+        y = self.weight[x]
+        yq = mx.dequantize(
+            w[x],
+            scales=s[x],
+            biases=b[x],
+            group_size=self.group_size,
+            bits=self.bits,
+        )
+        return (y - mx.stop_gradient(y)) + mx.stop_gradient(yq)
+
+    def as_linear(self, x):
+        # Quantize and then matmul
+        w, s, b = mx.quantize(self.weight, self.group_size, self.bits)
+        y = x @ self.weight.T
+        yq = mx.quantized_matmul(
+            x,
+            w,
+            scales=s,
+            biases=b,
+            transpose=True,
+            group_size=self.group_size,
+            bits=self.bits,
+        )
+        return (y - mx.stop_gradient(y)) + mx.stop_gradient(yq)
+
+    @classmethod
+    def from_embedding(
+        cls, embedding_layer: nn.Module, group_size: int = 64, bits: int = 4
+    ):
+        embedding_dims, dims = embedding_layer.weight.shape
+        ql = cls(embedding_dims, dims, group_size, bits)
+        ql.weight = embedding_layer.weight
+        return ql
+
+
+class StraightThroughQuantizedLinear(nn.Module):
+    def __init__(
+        self,
+        input_dims: int,
+        output_dims: int,
+        bias: bool = True,
+        group_size: int = 64,
+        bits: int = 4,
+    ):
+        super().__init__()
+
+        # Quantization config
+        self.group_size = group_size
+        self.bits = bits
+
+        self.weight = mx.zeros(shape=(output_dims, input_dims))
+        if bias:
+            self.bias = mx.zeros((output_dims,))
+
+    def __call__(self, x):
+        # Quantize and then matmul
+        w, s, b = mx.quantize(self.weight, self.group_size, self.bits)
+        y = x @ self.weight.T
+        yq = mx.quantized_matmul(
+            x,
+            w,
+            scales=s,
+            biases=b,
+            transpose=True,
+            group_size=self.group_size,
+            bits=self.bits,
+        )
+        x = (y - mx.stop_gradient(y)) + mx.stop_gradient(yq)
+        if "bias" in self:
+            x = x + self["bias"]
+        return x
+
+    @classmethod
+    def from_linear(cls, linear_layer: nn.Module, group_size: int = 64, bits: int = 4):
+        output_dims, input_dims = linear_layer.weight.shape
+        ql = cls(input_dims, output_dims, False, group_size, bits)
+        if "bias" in linear_layer:
+            ql.bias = linear_layer.bias
+        return ql
+
+
+def quantize(
+    model: nn.Module,
+    group_size: int = 64,
+    bits: int = 4,
+):
+    def _maybe_quantize(path, m):
+        if isinstance(m, nn.Linear):
+            return StraightThroughQuantizedLinear.from_linear(
+                m, group_size=group_size, bits=bits
+            )
+        elif isinstance(m, nn.Embedding):
+            return StraightThroughQuantizedEmbedding.from_embedding(
+                m, group_size=group_size, bits=bits
+            )
+        else:
+            return m
+
+    leaves = tree_map_with_path(
+        _maybe_quantize, model.leaf_modules(), is_leaf=nn.Module.is_module
+    )
+    model.update_modules(leaves)
+
+
+def qat(
+    model,
+    opt,
+    data,
+    group_size: int = 64,
+    bits: int = 3,
+    batch_size: int = 2,
+    max_seq_length: int = 2048,
+    temperature: float = 0.5,
+    dtype: mx.Dtype = mx.bfloat16,
+):
+    group = mx.distributed.init()
+    world_size = group.size()
+    rank = group.rank()
+
+    def log_norm(x):
+        x = x * (1 / temperature)
+        return x - mx.logsumexp(x, axis=-1, keepdims=True)
+
+    q_model = copy.deepcopy(model)
+    quantize(q_model, bits=bits, group_size=group_size)
+
+    def loss_fn(params, x, targets, lengths):
+        q_model.update(tree_map(lambda x: x.astype(dtype), params))
+        logits = q_model(x).astype(mx.float32)
+        losses = nn.losses.kl_div_loss(log_norm(logits), targets, reduction="none")
+        mask = mx.arange(targets.shape[1]) < lengths[:, 1:]
+        ntoks = mask.sum()
+        loss = (mask * losses).sum() / ntoks
+        return loss, ntoks
+
+    def step(inputs, targets, lengths, params):
+        (loss, ntoks), grads = mx.value_and_grad(loss_fn)(
+            params, inputs, targets, lengths
+        )
+        grads = nn.average_gradients(grads)
+        params = opt.apply_gradients(grads, params)
+        return loss, ntoks, params
+
+    # Accumulate learned weights in higher precision
+    params = tree_map(
+        lambda x: x.astype(mx.float32),
+        model.trainable_parameters(),
+    )
+
+    avg_loss = None
+    tokens = 0
+    tic = time.time()
+    for it, (batch, lengths) in enumerate(
+        iterate_batches(data, batch_size, max_seq_length)
+    ):
+        targets = log_norm(model(batch).astype(mx.float32))
+        mx.eval(targets)
+        loss, ntoks, params = step(batch, targets, lengths, params)
+        mx.eval(loss, params)
+        loss = mx.distributed.all_sum(loss, stream=mx.cpu).item() / world_size
+        ntoks = mx.distributed.all_sum(ntoks, stream=mx.cpu).item()
+        tokens += ntoks
+        toks_per_sec = tokens / (time.time() - tic)
+        avg_loss = 0.95 * (avg_loss or loss) + 0.05 * loss
+        if rank == 0:
+            print(
+                f"{it=}, {loss=:.3f}, {avg_loss=:.4f}, {tokens=}, {toks_per_sec=:.3f}",
+                flush=True,
+            )
+    model.update(tree_map(lambda x: x.astype(dtype), params))
+
+
+def save_model(
+    model: nn.Module,
+    tokenizer: TokenizerWrapper,
+    config,
+    model_path: Path,
+    mlx_path: str,
+    hf_path: str,
+):
+    weights = dict(tree_flatten(model.parameters()))
+
+    mlx_path = Path(mlx_path)
+    save_weights(mlx_path, weights, donate_weights=True)
+
+    py_files = glob.glob(str(model_path / "*.py"))
+    for file in py_files:
+        shutil.copy(file, mlx_path)
+
+    tokenizer.save_pretrained(mlx_path)
+
+    save_config(config, config_path=mlx_path / "config.json")
+    create_model_card(mlx_path, hf_path)
+
+
+def load_data(tokenizer, data_path: str, num_samples: int):
+    args = types.SimpleNamespace(
+        hf_dataset={
+            "path": data_path,
+            "train_split": f"train[:{num_samples}]",
+            "valid_split": "train[:1]",
+        },
+        train=True,
+        test=False,
+    )
+    dataset = load_dataset(args, tokenizer)[0]
+    return [dataset.process(d) for d in dataset]
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", "-m", default="Qwen/Qwen3-1.7B")
+    parser.add_argument(
+        "--mlx-path", default="mlx_model", help="Path to save the quantized model."
+    )
+    parser.add_argument(
+        "--bits",
+        type=int,
+        default=4,
+        help="Bits per weight for quantization.",
+    )
+    parser.add_argument(
+        "--group-size", type=int, default=64, help="Group size for quantization."
+    )
+    parser.add_argument(
+        "--num-samples",
+        type=int,
+        default=1024,
+        help="Number of samples to use for training.",
+    )
+    parser.add_argument("--max-seq-length", type=int, default=2048)
+    parser.add_argument("--seed", type=int, default=123)
+    parser.add_argument("--learning-rate", type=float, default=1e-5)
+    parser.add_argument("--batch-size", type=int, default=8)
+    parser.add_argument(
+        "--data-path",
+        type=str,
+        default="allenai/tulu-3-sft-mixture",
+        help="A Hugging Face dataset which is compatible with an mlx-lm dataset format.",
+    )
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=0.5,
+        help="Temperature scaling for the loss.",
+    )
+    args = parser.parse_args()
+
+    group = mx.distributed.init()
+
+    num_samples = args.num_samples
+    if num_samples % group.size() > 0:
+        num_samples += group.size() - num_samples % group.size()
+
+    np.random.seed(args.seed)
+    mx.random.seed(args.seed)
+
+    model_path = get_model_path(args.model, revision=None)
+    model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
+    if "quantization" in config:
+        raise ValueError("Teacher model for QAT training should not be quantized")
+
+    calibration_data = load_data(tokenizer, args.data_path, args.num_samples)
+
+    q_model = copy.deepcopy(model)
+    tree_flatten(q_model.parameters())
+    opt = optimizers.Adam(learning_rate=args.learning_rate, bias_correction=True)
+    qat(
+        model,
+        opt,
+        calibration_data,
+        bits=args.bits,
+        group_size=args.group_size,
+        batch_size=args.batch_size,
+        max_seq_length=args.max_seq_length,
+        temperature=args.temperature,
+    )
+    _, config = quantize_model(
+        model,
+        config,
+        q_group_size=args.group_size,
+        q_bits=args.bits,
+    )
+    save_model(model, tokenizer, config, model_path, args.mlx_path, args.model)
+
+
+if __name__ == "__main__":
+    main()

mlx_lm/sample_utils.py

@@ -2,7 +2,7 @@
 
 import math
 from functools import partial
-from typing import Callable, Dict, Optional
+from typing import Callable, Dict, List, Optional
 
 import mlx.core as mx
 
@@ -12,7 +12,10 @@ def make_sampler(
     top_p: float = 0.0,
     min_p: float = 0.0,
     min_tokens_to_keep: int = 1,
-    top_k: int = -1,
+    top_k: int = 0,
+    xtc_probability: float = 0.0,
+    xtc_threshold: float = 0.0,
+    xtc_special_tokens: List[int] = [],
 ) -> Callable[mx.array, mx.array]:
     """
     Make a sampler function for use with ``generate_step``.
@@ -28,6 +31,13 @@ def make_sampler(
           be filtered by min_p sampling.
         top_k (int, optional): The top k tokens ranked by probability to constrain
           the sampling to.
+        xtc_probability (float, optional): The probability of applying XTC
+            sampling.
+        xtc_threshold (float, optional): The threshold the probs need to reach
+            for being sampled.
+        xtc_special_tokens (list(int), optional): List of special tokens IDs to
+            be excluded from XTC sampling.
+
 
     Returns:
         Callable[mx.array, mx.array]:
@@ -44,6 +54,10 @@ def make_sampler(
         sampling_methods.append(lambda x: apply_top_p(x, top_p))
     if min_p != 0.0:
         sampling_methods.append(lambda x: apply_min_p(x, min_p, min_tokens_to_keep))
+    if xtc_probability > 0.0:
+        sampling_methods.append(
+            lambda x: apply_xtc(x, xtc_probability, xtc_threshold, xtc_special_tokens)
+        )
 
     # Apply the sampling methods
     def sampler(logits):
@@ -194,7 +208,6 @@ def apply_top_p(logits: mx.array, top_p: float) -> mx.array:
     """
     # referenced implementation from https://github.com/huggingface/transformers/blob/main/src/transformers/generation/logits_process.py#L449-L460
     probs = mx.softmax(logits, axis=-1)
-
     # sort probs in ascending order
     sorted_indices = mx.argsort(probs, axis=-1)
     sorted_probs = mx.take_along_axis(probs, sorted_indices, axis=-1)
@@ -219,6 +232,43 @@ def apply_top_p(logits: mx.array, top_p: float) -> mx.array:
     return mx.log(original_order_probs)
 
 
+@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
+def apply_xtc(
+    logits: mx.array,
+    xtc_probability: float,
+    xtc_threshold: float,
+    xtc_special_tokens: List[int],
+) -> mx.array:
+    """
+    Apply XTC sampling to the logits.
+
+    Args:
+        logits: The logits from the model's output.
+        xtc_probability (float): Probability of XTC sampling to happen for each token
+        xtc_threshold (float): The threshold the probs need to reach for being sampled.
+        special_tokens_ids (list(int)): List of special tokens IDs to be excluded from XTC sampling.
+    """
+    if not (0 <= xtc_threshold <= 0.5):
+        raise ValueError(
+            f"`threshold` has to be a float in the [0, 0.5] interval, but is {xtc_threshold}"
+        )
+    if not (0 <= xtc_probability <= 1.0):
+        raise ValueError(
+            f"`probability` has to be a float in the [0, 1] interval, but is {xtc_probability}"
+        )
+
+    probs = mx.softmax(logits, -1)
+    mask = probs > mx.where(probs > xtc_threshold, probs, mx.inf).min()
+    if xtc_special_tokens:
+        mask[..., xtc_special_tokens] = False
+
+    return mx.where(
+        mx.random.uniform(0, 1) > xtc_probability,
+        logits,
+        mx.where(mask, -mx.inf, logits),
+    )
+
+
 @partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
 def categorical_sampling(logits, temp):
     return mx.random.categorical(logits * (1 / temp))

mlx_lm/server.py

@@ -30,7 +30,7 @@ from ._version import __version__
 from .generate import stream_generate
 from .models.cache import can_trim_prompt_cache, make_prompt_cache, trim_prompt_cache
 from .sample_utils import make_logits_processors, make_sampler
-from .utils import load
+from .utils import common_prefix_len, load
 
 
 def get_system_fingerprint():
@@ -146,7 +146,7 @@ def process_message_content(messages):
 @dataclass
 class PromptCache:
     cache: List[Any] = field(default_factory=list)
-    model_key: Tuple[str, Optional[str]] = ("", None)
+    model_key: Tuple[str, Optional[str]] = ("", None, None)
     tokens: List[int] = field(default_factory=list)
 
 
@@ -157,10 +157,11 @@ class ModelProvider:
         self.model_key = None
         self.model = None
         self.tokenizer = None
+        self.draft_model = None
 
         # Preload the default model if it is provided
         if self.cli_args.model is not None:
-            self.load("default_model")
+            self.load("default_model", draft_model_path="default_model")
 
     def _validate_model_path(self, model_path: str):
         model_path = Path(model_path)
@@ -170,14 +171,15 @@ class ModelProvider:
             )
 
     # Added in adapter_path to load dynamically
-    def load(self, model_path, adapter_path=None):
-        if self.model_key == (model_path, adapter_path):
+    def load(self, model_path, adapter_path=None, draft_model_path=None):
+        if self.model_key == (model_path, adapter_path, draft_model_path):
             return self.model, self.tokenizer
 
         # Remove the old model if it exists.
         self.model = None
         self.tokenizer = None
         self.model_key = None
+        self.draft_model = None
 
         # Building tokenizer_config
         tokenizer_config = {
@@ -186,7 +188,12 @@ class ModelProvider:
         if self.cli_args.chat_template:
             tokenizer_config["chat_template"] = self.cli_args.chat_template
 
-        if model_path == "default_model" and self.cli_args.model is not None:
+        if model_path == "default_model":
+            if self.cli_args.model is None:
+                raise ValueError(
+                    "A model path has to be given as a CLI "
+                    "argument or in the HTTP request"
+                )
             model, tokenizer = load(
                 self.cli_args.model,
                 adapter_path=(
@@ -204,10 +211,30 @@ class ModelProvider:
             if tokenizer.chat_template is None:
                 tokenizer.chat_template = tokenizer.default_chat_template
 
-        self.model_key = (model_path, adapter_path)
+        self.model_key = (model_path, adapter_path, draft_model_path)
         self.model = model
         self.tokenizer = tokenizer
 
+        def validate_draft_tokenizer(draft_tokenizer):
+            # Check if tokenizers are compatible
+            if draft_tokenizer.vocab_size != tokenizer.vocab_size:
+                logging.warning(
+                    "Draft model tokenizer does not match model tokenizer. "
+                    "Speculative decoding may not work as expected."
+                )
+
+        # Load draft model if specified
+        if (
+            draft_model_path == "default_model"
+            and self.cli_args.draft_model is not None
+        ):
+            self.draft_model, draft_tokenizer = load(self.cli_args.draft_model)
+            validate_draft_tokenizer(draft_tokenizer)
+
+        elif draft_model_path is not None and draft_model_path != "default_model":
+            self._validate_model_path(draft_model_path)
+            self.draft_model, draft_tokenizer = load(draft_model_path)
+            validate_draft_tokenizer(draft_tokenizer)
         return self.model, self.tokenizer
 
 
@@ -279,6 +306,8 @@ class APIHandler(BaseHTTPRequestHandler):
         self.stream = self.body.get("stream", False)
         self.stream_options = self.body.get("stream_options", None)
         self.requested_model = self.body.get("model", "default_model")
+        self.requested_draft_model = self.body.get("draft_model", "default_model")
+        self.num_draft_tokens = self.body.get("num_draft_tokens", 3)
         self.adapter = self.body.get("adapters", None)
         self.max_tokens = self.body.get("max_completion_tokens", None)
         if self.max_tokens is None:
@@ -287,14 +316,17 @@ class APIHandler(BaseHTTPRequestHandler):
         self.top_p = self.body.get("top_p", 1.0)
         self.repetition_penalty = self.body.get("repetition_penalty", 1.0)
         self.repetition_context_size = self.body.get("repetition_context_size", 20)
+        self.xtc_probability = self.body.get("xtc_probability", 0.0)
+        self.xtc_threshold = self.body.get("xtc_threshold", 0.0)
         self.logit_bias = self.body.get("logit_bias", None)
         self.logprobs = self.body.get("logprobs", -1)
         self.validate_model_parameters()
-
         # Load the model if needed
         try:
             self.model, self.tokenizer = self.model_provider.load(
-                self.requested_model, self.adapter
+                self.requested_model,
+                self.adapter,
+                self.requested_draft_model,
             )
         except:
             self._set_completion_headers(404)
@@ -363,7 +395,15 @@ class APIHandler(BaseHTTPRequestHandler):
                 self.logit_bias = {int(k): v for k, v in self.logit_bias.items()}
             except ValueError:
                 raise ValueError("logit_bias must be a dict of int to float")
-
+        if not (
+            isinstance(self.xtc_probability, float)
+            and 0.00 <= self.xtc_probability <= 1.00
+        ):
+            raise ValueError(f"xtc_probability must be a float between 0.00 and 1.00")
+        if not (
+            isinstance(self.xtc_threshold, float) and 0.00 <= self.xtc_threshold <= 0.50
+        ):
+            raise ValueError(f"xtc_threshold must be a float between 0.00 and 0.5")
         if not isinstance(self.requested_model, str):
             raise ValueError("model must be a string")
         if self.adapter is not None and not isinstance(self.adapter, str):
@@ -452,34 +492,84 @@ class APIHandler(BaseHTTPRequestHandler):
 
         return response
 
+    def reset_prompt_cache(self, prompt):
+        """Resets the prompt cache and associated state.
+
+        Args:
+            prompt (List[int]): The tokenized new prompt which will populate the
+                reset cache.
+        """
+        logging.debug(f"*** Resetting cache. ***")
+        self.prompt_cache.model_key = self.model_provider.model_key
+        self.prompt_cache.cache = make_prompt_cache(self.model_provider.model)
+        if self.model_provider.draft_model is not None:
+            self.prompt_cache.cache += make_prompt_cache(
+                self.model_provider.draft_model
+            )
+        self.prompt_cache.tokens = list(prompt)  # Cache the new prompt fully
+
     def get_prompt_cache(self, prompt):
+        """
+        Determines the portion of the prompt that needs processing by comparing
+        it to the cached prompt and attempting to reuse the common prefix.
+
+        This function updates the internal prompt cache state (tokens and model cache)
+        based on the comparison. If a common prefix exists, it attempts to trim
+        the model cache (if supported) to match the common prefix length, avoiding
+        recomputation.
+
+        Args:
+            prompt (List[int]): The tokenized new prompt.
+
+        Returns:
+            List[int]: The suffix of the prompt that actually needs to be processed
+                       by the model. This will be the full prompt if the cache is
+                       reset or cannot be effectively used.
+        """
         cache_len = len(self.prompt_cache.tokens)
         prompt_len = len(prompt)
-        prefix_len = min(cache_len, prompt_len)
+        com_prefix_len = common_prefix_len(self.prompt_cache.tokens, prompt)
+
+        # Condition 1: Model changed or no common prefix at all. Reset cache.
         if (
             self.prompt_cache.model_key != self.model_provider.model_key
-            or prompt[:prefix_len] != self.prompt_cache.tokens[:prefix_len]
+            or com_prefix_len == 0
         ):
-            self.prompt_cache.model_key = self.model_provider.model_key
-            self.prompt_cache.cache = make_prompt_cache(self.model_provider.model)
-            self.prompt_cache.tokens = []
-        elif cache_len >= prompt_len:
-            # Trim the cache if it contains the prompt as a prefix. This case
-            # is useful for reusing the cache for multiple queries with a long
-            # prompt
+            self.reset_prompt_cache(prompt)
+
+        # Condition 2: Common prefix exists and matches cache length. Process suffix.
+        elif com_prefix_len == cache_len:
+            logging.debug(
+                f"*** Cache is prefix of prompt (cache_len: {cache_len}, prompt_len: {prompt_len}). Processing suffix. ***"
+            )
+            prompt = prompt[com_prefix_len:]
+            self.prompt_cache.tokens.extend(prompt)
+
+        # Condition 3: Common prefix exists but is shorter than cache length. Attempt trim.
+        elif com_prefix_len < cache_len:
+            logging.debug(
+                f"*** Common prefix ({com_prefix_len}) shorter than cache ({cache_len}). Attempting trim. ***"
+            )
+
             if can_trim_prompt_cache(self.prompt_cache.cache):
-                num_to_trim = cache_len - prompt_len + 1
+                num_to_trim = cache_len - com_prefix_len
+                logging.debug(f"    Trimming {num_to_trim} tokens from cache.")
                 trim_prompt_cache(self.prompt_cache.cache, num_to_trim)
-                self.prompt_cache.tokens = self.prompt_cache.tokens[:-num_to_trim]
-                prompt = prompt[-1:]
+                self.prompt_cache.tokens = self.prompt_cache.tokens[:com_prefix_len]
+                prompt = prompt[com_prefix_len:]
+                self.prompt_cache.tokens.extend(prompt)
             else:
-                self.prompt_cache.cache = make_prompt_cache(self.model_provider.model)
-                self.prompt_cache.tokens = []
+                logging.debug(f"    Cache cannot be trimmed. Resetting cache.")
+                self.reset_prompt_cache(prompt)
+
+        # This case should logically not be reached if com_prefix_len <= cache_len
         else:
-            # Trim the prompt if it contains the cache as a prefix. This case
-            # is to avoid recomputing the cache in multi-turn chats.
-            prompt = prompt[cache_len:]
-        self.prompt_cache.tokens.extend(prompt)
+            logging.error(
+                f"Unexpected cache state: com_prefix_len ({com_prefix_len}) > cache_len ({cache_len}). Resetting cache."
+            )
+            self.reset_prompt_cache(prompt)
+
+        logging.debug(f"Returning {len(prompt)} tokens for processing.")
         return prompt
 
     def handle_completion(
@@ -510,10 +600,22 @@ class APIHandler(BaseHTTPRequestHandler):
 
         text = ""
         tic = time.perf_counter()
-        sampler = make_sampler(self.temperature, top_p=self.top_p)
-        logits_processors = make_logits_processors(
-            self.logit_bias, self.repetition_penalty, self.repetition_context_size
+        sampler = make_sampler(
+            self.temperature,
+            top_p=self.top_p,
+            xtc_probability=self.xtc_probability,
+            xtc_threshold=self.xtc_threshold,
+            xtc_special_tokens=[
+                self.tokenizer.eos_token_id,
+                self.tokenizer.encode("\n"),
+            ],
         )
+        logits_processors = make_logits_processors(
+            self.logit_bias,
+            self.repetition_penalty,
+            self.repetition_context_size,
+        )
+
         for gen_response in stream_generate(
             model=self.model,
             tokenizer=self.tokenizer,
@@ -522,6 +624,8 @@ class APIHandler(BaseHTTPRequestHandler):
             sampler=sampler,
             logits_processors=logits_processors,
             prompt_cache=self.prompt_cache.cache,
+            draft_model=self.model_provider.draft_model,
+            num_draft_tokens=self.num_draft_tokens,
         ):
             segment = gen_response.text
             text += segment
@@ -680,10 +784,20 @@ class APIHandler(BaseHTTPRequestHandler):
         self._set_completion_headers(200)
         self.end_headers()
 
+        files = ["config.json", "model.safetensors.index.json", "tokenizer_config.json"]
+
+        def probably_mlx_lm(repo):
+            if repo.repo_type != "model":
+                return False
+            if "main" not in repo.refs:
+                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)
+
         # Scan the cache directory for downloaded mlx models
         hf_cache_info = scan_cache_dir()
         downloaded_models = [
-            repo for repo in hf_cache_info.repos if "mlx" in repo.repo_id
+            repo for repo in hf_cache_info.repos if probably_mlx_lm(repo)
         ]
 
         # Create a list of available models
@@ -758,6 +872,12 @@ def main():
         default=8080,
         help="Port for the HTTP server (default: 8080)",
     )
+    parser.add_argument(
+        "--draft-model",
+        type=str,
+        help="A model to be used for speculative decoding.",
+        default=None,
+    )
     parser.add_argument(
         "--trust-remote-code",
         action="store_true",
@@ -792,4 +912,8 @@ def main():
 
 
 if __name__ == "__main__":
+    print(
+        "Calling `python -m mlx_lm.server...` directly is deprecated."
+        " Use `mlx_lm.server...` or `python -m mlx_lm server ...` instead."
+    )
     main()

mlx_lm/tokenizer_utils.py

@@ -1,5 +1,6 @@
 import json
 from functools import partial
+from json import JSONDecodeError
 from typing import List
 
 from transformers import AutoTokenizer
@@ -341,7 +342,9 @@ def _is_bpe_decoder(decoder):
     return isinstance(decoder, dict) and decoder.get("type", None) == "ByteLevel"
 
 
-def load_tokenizer(model_path, tokenizer_config_extra={}, eos_token_ids=None):
+def load_tokenizer(
+    model_path, tokenizer_config_extra={}, return_tokenizer=True, eos_token_ids=None
+):
     """Load a huggingface tokenizer and try to infer the type of streaming
     detokenizer to use.
 
@@ -353,7 +356,11 @@ def load_tokenizer(model_path, tokenizer_config_extra={}, eos_token_ids=None):
     tokenizer_file = model_path / "tokenizer.json"
     if tokenizer_file.exists():
         with open(tokenizer_file, "r", encoding="utf-8") as fid:
-            tokenizer_content = json.load(fid)
+            try:
+                tokenizer_content = json.load(fid)
+            except JSONDecodeError as e:
+                raise JSONDecodeError("Failed to parse tokenizer.json", e.doc, e.pos)
+
         if "decoder" in tokenizer_content:
             if _is_spm_decoder(tokenizer_content["decoder"]):
                 detokenizer_class = SPMStreamingDetokenizer
@@ -364,11 +371,15 @@ def load_tokenizer(model_path, tokenizer_config_extra={}, eos_token_ids=None):
 
     if isinstance(eos_token_ids, int):
         eos_token_ids = [eos_token_ids]
-    return TokenizerWrapper(
-        AutoTokenizer.from_pretrained(model_path, **tokenizer_config_extra),
-        detokenizer_class,
-        eos_token_ids=eos_token_ids,
-    )
+
+    if return_tokenizer:
+        return TokenizerWrapper(
+            AutoTokenizer.from_pretrained(model_path, **tokenizer_config_extra),
+            detokenizer_class,
+            eos_token_ids=eos_token_ids,
+        )
+    else:
+        return detokenizer_class
 
 
 def no_bos_or_eos(sequence: List, bos: int, eos: int) -> List:

mlx_lm/tuner/datasets.py

@@ -17,7 +17,7 @@ class TextDataset:
         tokenizer: PreTrainedTokenizer,
         text_key: str = "text",
     ):
-        self._data = [d for d in data]
+        self._data = data
         self.tokenizer = tokenizer
         self.text_key = text_key
 
@@ -47,7 +47,7 @@ class ChatDataset:
         chat_key: str = "messages",
         mask_prompt: bool = False,
     ):
-        self._data = [d for d in data]
+        self._data = data
         self.chat_key = chat_key
         self.mask_prompt = mask_prompt
         self.tokenizer = tokenizer
@@ -58,14 +58,11 @@ class ChatDataset:
         tokens = self.tokenizer.apply_chat_template(messages, tools=tools)
         if self.mask_prompt:
             messages = messages[:-1]
-            offset = len(tokenizer.apply_chat_template(messages, tools=tools))
+            offset = len(self.tokenizer.apply_chat_template(messages, tools=tools))
             return (tokens, offset)
         else:
             return tokens
 
-    def itemlen(idx: int):
-        return len(self._data[idx])
-
     def __getitem__(self, idx: int):
         return self._data[idx]
 
@@ -88,7 +85,7 @@ class CompletionsDataset:
         completion_key: str,
         mask_prompt: bool,
     ):
-        self._data = [d for d in data]
+        self._data = data
         self.prompt_key = prompt_key
         self.completion_key = completion_key
         self.mask_prompt = mask_prompt
@@ -124,12 +121,17 @@ class ConcatenatedDataset:
         self._len = sum(len(d) for d in self._data)
 
     def __getitem__(self, idx: int):
-        for data in self._data:
+        for data_idx, data in enumerate(self._data):
             j = idx - len(data)
             if j < 0:
                 break
             idx = j
-        return data[idx]
+        datum = data[idx]
+        datum["_dataset"] = data_idx
+        return datum
+
+    def process(self, d):
+        return self._data[d["_dataset"]].process(d)
 
     def __len__(self):
         return self._len

mlx_lm/tuner/trainer.py

@@ -72,21 +72,19 @@ def default_loss(model, batch, lengths):
     targets = batch[:, 1:]
 
     logits = model(inputs)
-    logits = logits.astype(mx.float32)
 
     steps = mx.arange(1, targets.shape[1] + 1)
     mask = mx.logical_and(steps >= lengths[:, 0:1], steps <= lengths[:, 1:])
 
     ce = nn.losses.cross_entropy(logits, targets) * mask
     ntoks = mask.sum()
-    ce = ce.sum() / ntoks
+    ce = ce.astype(mx.float32).sum() / ntoks
 
     return ce, ntoks
 
 
 def iterate_batches(
     dataset,
-    tokenizer,
     batch_size,
     max_seq_length,
     train=False,
@@ -154,7 +152,6 @@ def iterate_batches(
 def evaluate(
     model,
     dataset,
-    tokenizer,
     batch_size,
     num_batches,
     max_seq_length=2048,
@@ -171,7 +168,6 @@ def evaluate(
         index_iterator,
         iterate_batches(
             dataset=dataset,
-            tokenizer=tokenizer,
             batch_size=batch_size,
             max_seq_length=max_seq_length,
         ),
@@ -200,7 +196,6 @@ class TrainingCallback:
 
 def train(
     model,
-    tokenizer,
     optimizer,
     train_dataset,
     val_dataset,
@@ -235,9 +230,6 @@ def train(
 
         return lvalue, toks
 
-    train_dataset = CacheDataset(train_dataset)
-    val_dataset = CacheDataset(val_dataset)
-
     loss_value_and_grad = nn.value_and_grad(model, loss)
 
     model.train()
@@ -251,7 +243,6 @@ def train(
         range(1, args.iters + 1),
         iterate_batches(
             dataset=train_dataset,
-            tokenizer=tokenizer,
             batch_size=args.batch_size,
             max_seq_length=args.max_seq_length,
             train=True,
@@ -266,7 +257,6 @@ def train(
                 model=model,
                 dataset=val_dataset,
                 loss=loss,
-                tokenizer=tokenizer,
                 batch_size=args.batch_size,
                 num_batches=args.val_batches,
                 max_seq_length=args.max_seq_length,

mlx_lm/tuner/utils.py

@@ -87,6 +87,8 @@ def linear_to_lora_layers(
         "hunyuan",
         "qwen2",
         "qwen2_moe",
+        "qwen3",
+        "qwen3_moe",
         "phimoe",
         "gemma",
         "gemma2",
@@ -103,6 +105,8 @@ def linear_to_lora_layers(
         "olmo2",
         "olmoe",
         "internlm3",
+        "glm4",
+        "mimo",
     ]:
         keys = set(["self_attn.q_proj", "self_attn.v_proj"])
         if model.model_type in ["mixtral", "phimoe"]:
@@ -110,7 +114,7 @@ def linear_to_lora_layers(
         if model.model_type == "qwen2_moe":
             keys.add("mlp.gate")
             keys.add("mlp.shared_expert_gate")
-        if model.model_type == "olmoe":
+        if model.model_type in ["olmoe", "qwen3_moe"]:
             keys.add("mlp.gate")
 
     elif model.model_type == "gpt_bigcode":

mlx_lm/upload.py

@@ -0,0 +1,22 @@
+# Copyright © 2025 Apple Inc.
+
+import argparse
+
+from .utils import upload_to_hub
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Upload a model to the Hugging Face Hub"
+    )
+
+    parser.add_argument(
+        "--path", type=str, default="mlx_model", help="Path to the MLX model."
+    )
+    parser.add_argument(
+        "--upload-repo",
+        help="The Hugging Face repo to upload the model to.",
+        type=str,
+    )
+    args = parser.parse_args()
+    upload_to_hub(args.path, args.upload_repo)

mlx_lm/utils.py

@@ -6,6 +6,7 @@ import importlib
 import json
 import logging
 import os
+import shutil
 from pathlib import Path
 from textwrap import dedent
 from typing import (
@@ -298,20 +299,15 @@ def make_shards(weights: dict, max_file_size_gb: int = MAX_FILE_SIZE_GB) -> list
     return shards
 
 
-def upload_to_hub(path: str, upload_repo: str, hf_path: str):
+def create_model_card(path: Union[str, Path], hf_path: Union[str, Path]):
     """
     Uploads the model to Hugging Face hub.
 
     Args:
-        path (str): Local path to the model.
-        upload_repo (str): Name of the HF repo to upload to.
-        hf_path (str): Path to the original Hugging Face model.
+        path (Union[str, Path]): Local path to the model.
+        hf_path (Union[str, Path]): Path to the original Hugging Face model.
     """
-    import os
-
-    from huggingface_hub import HfApi, ModelCard, logging
-
-    from . import __version__
+    from huggingface_hub import ModelCard
 
     card = ModelCard.load(hf_path)
     card.data.library_name = "mlx"
@@ -320,7 +316,27 @@ def upload_to_hub(path: str, upload_repo: str, hf_path: str):
         card.data.tags = ["mlx"]
     elif "mlx" not in card.data.tags:
         card.data.tags += ["mlx"]
-    card.data.base_model = hf_path
+    card.data.base_model = str(hf_path)
+    card.text = ""
+    card.save(os.path.join(path, "README.md"))
+
+
+def upload_to_hub(path: str, upload_repo: str):
+    """
+    Uploads the model to Hugging Face hub.
+
+    Args:
+        path (str): Local path to the model.
+        upload_repo (str): Name of the HF repo to upload to.
+    """
+    from huggingface_hub import HfApi, ModelCard, logging
+
+    from . import __version__
+
+    logging.set_verbosity_info()
+    card_path = Path(path) / "README.md"
+    card = ModelCard.load(card_path)
+    hf_path = card.data.base_model
     card.text = dedent(
         f"""
         # {upload_repo}
@@ -352,9 +368,7 @@ def upload_to_hub(path: str, upload_repo: str, hf_path: str):
         ```
         """
     )
-    card.save(os.path.join(path, "README.md"))
-
-    logging.set_verbosity_info()
+    card.save(card_path)
 
     api = HfApi()
     api.create_repo(repo_id=upload_repo, exist_ok=True)
@@ -491,3 +505,52 @@ def save_config(
     # write the updated config to the config_path (if provided)
     with open(config_path, "w") as fid:
         json.dump(config, fid, indent=4)
+
+
+def save(
+    dst_path: Union[str, Path],
+    src_path: Union[str, Path],
+    weights: Dict[str, mx.array],
+    tokenizer: TokenizerWrapper,
+    config: Dict[str, Any],
+    hf_repo: Optional[str] = None,
+    donate_weights: bool = True,
+):
+    src_path = Path(src_path)
+    dst_path = Path(dst_path)
+    save_weights(dst_path, weights, donate_weights=True)
+    save_config(config, config_path=dst_path / "config.json")
+    tokenizer.save_pretrained(dst_path)
+
+    for p in ["*.py", "generation_config.json"]:
+        for file in glob.glob(str(src_path / p)):
+            shutil.copy(file, dst_path)
+
+    if hf_repo is not None:
+        create_model_card(dst_path, hf_repo)
+
+
+def common_prefix_len(list1, list2):
+    """
+    Calculates the length of the common prefix of two lists.
+
+    Args:
+        list1: The first list of strings.
+        list2: The second list of strings.
+
+    Returns:
+        The length of the common prefix. Returns 0 if lists are empty
+        or do not match at the first element.
+    """
+    # Determine the maximum possible length of the common prefix
+    min_len = min(len(list1), len(list2))
+
+    # Iterate up to the length of the shorter list
+    for i in range(min_len):
+        if list1[i] != list2[i]:
+            # Mismatch found, the common prefix length is the current index
+            return i
+
+    # No mismatch found within the bounds of the shorter list,
+    # so the common prefix length is the length of the shorter list.
+    return min_len

requirements.txt

@@ -1,4 +1,4 @@
-mlx>=0.24.1
+mlx>=0.25.0
 numpy
 transformers[sentencepiece]>=4.39.3
 protobuf

setup.py

@@ -29,9 +29,12 @@ setup(
     extras_require={
         "test": ["datasets"],
         "evaluate": ["lm-eval", "tqdm"],
+        "lwq": ["datasets"],
     },
     entry_points={
         "console_scripts": [
+            "mlx_lm.awq = mlx_lm.awq:main",
+            "mlx_lm.dwq = mlx_lm.dwq:main",
             "mlx_lm.cache_prompt = mlx_lm.cache_prompt:main",
             "mlx_lm.chat = mlx_lm.chat:main",
             "mlx_lm.convert = mlx_lm.convert:main",
@@ -42,6 +45,7 @@ setup(
             "mlx_lm.merge = mlx_lm.merge:main",
             "mlx_lm.server = mlx_lm.server:main",
             "mlx_lm.manage = mlx_lm.manage:main",
+            "mlx_lm.upload = mlx_lm.upload:main",
         ]
     },
 )

tests/test_finetune.py

@@ -67,7 +67,7 @@ class TestLora(unittest.TestCase):
             )
             self.assertEqual(trainable_params, expected_trainable_parameters)
 
-        params = {"rank": 8, "alpha": 16, "dropout": 0.0, "scale": 10.0}
+        params = {"rank": 8, "dropout": 0.0, "scale": 10.0}
         check_config(params)
 
         params["rank"] = 1
@@ -108,7 +108,7 @@ class TestLora(unittest.TestCase):
         )
 
         num_lora_layers = 4
-        params = {"rank": 8, "alpha": 16, "dropout": 0.0, "scale": 10.0}
+        params = {"rank": 8, "dropout": 0.0, "scale": 10.0}
 
         model = gpt_neox.Model(args)
         model.freeze()
@@ -365,7 +365,6 @@ class TestScheduleConfig(unittest.TestCase):
     def test_evaluate_calls(self):
         mock_model = MagicMock()
         mock_dataset = MagicMock()
-        mock_tokenizer = MagicMock()
         mock_default_loss = MagicMock()
         mock_iterate_batches = MagicMock()
 
@@ -388,7 +387,6 @@ class TestScheduleConfig(unittest.TestCase):
             evaluate(
                 model=mock_model,
                 dataset=mock_dataset,
-                tokenizer=mock_tokenizer,
                 batch_size=2,
                 num_batches=2,
                 max_seq_length=2048,
@@ -398,7 +396,6 @@ class TestScheduleConfig(unittest.TestCase):
 
         mock_iterate_batches.assert_called_once_with(
             dataset=mock_dataset,
-            tokenizer=mock_tokenizer,
             batch_size=2,
             max_seq_length=2048,
         )
@@ -407,7 +404,6 @@ class TestScheduleConfig(unittest.TestCase):
     def test_evaluate_infinite_batches(self):
         mock_model = MagicMock()
         mock_dataset = MagicMock()
-        mock_tokenizer = MagicMock()
         mock_default_loss = MagicMock()
         mock_iterate_batches = MagicMock()
 
@@ -427,7 +423,6 @@ class TestScheduleConfig(unittest.TestCase):
             evaluate(
                 model=mock_model,
                 dataset=mock_dataset,
-                tokenizer=mock_tokenizer,
                 batch_size=2,
                 num_batches=-1,
                 max_seq_length=2048,
@@ -437,7 +432,6 @@ class TestScheduleConfig(unittest.TestCase):
 
         mock_iterate_batches.assert_called_once_with(
             dataset=mock_dataset,
-            tokenizer=mock_tokenizer,
             batch_size=2,
             max_seq_length=2048,
         )

tests/test_generate.py

@@ -66,11 +66,15 @@ class TestGenerate(unittest.TestCase):
 
         # make a determinate sampler
         sampler = make_sampler(temp=0.0)
+        messages = [{"role": "user", "content": "hello"}]
+        prompt = self.tokenizer.apply_chat_template(
+            messages, add_generation_prompt=True
+        )
 
         for generation_result in stream_generate(
             model=self.model,
             tokenizer=self.tokenizer,
-            prompt="hello",
+            prompt=prompt,
             max_tokens=5,
             draft_model=draft_model,
             num_draft_tokens=2,
@@ -79,7 +83,8 @@ class TestGenerate(unittest.TestCase):
             drafted.append(generation_result.from_draft)
             results.append(generation_result)
 
-        self.assertEqual(len(results), 5)
+        self.assertEqual(len(results), 6)
+        drafted.pop()
         # since num_draft_tokens is 2 and draft model is the same, the
         # first 2 generations should be drafts, the third should come
         # from the target model, and last two should be drafts

tests/test_models.py

@@ -6,7 +6,7 @@ import mlx.nn as nn
 from mlx.utils import tree_map
 
 from mlx_lm.models import rope_utils
-from mlx_lm.models.base import create_causal_mask
+from mlx_lm.models.base import create_causal_mask, scaled_dot_product_attention
 from mlx_lm.models.cache import KVCache, RotatingKVCache, make_prompt_cache
 
 
@@ -166,6 +166,42 @@ class TestModels(unittest.TestCase):
         )
         self.assertTrue(isinstance(rope, rope_utils.Llama3RoPE))
 
+    def test_quantized_sdpa(self):
+        cache = KVCache()
+
+        k = 1e-1 * mx.random.normal(shape=(1, 1, 256, 32))
+        v = 1e-1 * mx.random.normal(shape=(1, 1, 256, 32))
+
+        cache.update_and_fetch(k, v)
+        quant_cache = cache.to_quantized(group_size=32, bits=8)
+
+        k = 1e-1 * mx.random.normal(shape=(1, 1, 1, 32))
+        v = 1e-1 * mx.random.normal(shape=(1, 1, 1, 32))
+
+        k_up, v_up = cache.update_and_fetch(k, v)
+        qk_up, qv_up = quant_cache.update_and_fetch(k, v)
+
+        q = 1e-1 * mx.random.normal(shape=(1, 4, 257, 32))
+
+        mask = "causal"
+        out = scaled_dot_product_attention(
+            q,
+            k_up,
+            v_up,
+            cache=cache,
+            mask=mask,
+            scale=1.0,
+        )
+        qout = scaled_dot_product_attention(
+            q,
+            qk_up,
+            qv_up,
+            cache=quant_cache,
+            mask=mask,
+            scale=1.0,
+        )
+        self.assertTrue(mx.allclose(out, qout, rtol=1e-2, atol=1e-2))
+
     def model_test_runner(self, model, model_type, vocab_size, num_layers):
 
         self.assertEqual(len(model.layers), num_layers)
@@ -307,6 +343,56 @@ class TestModels(unittest.TestCase):
             args.n_layers,
         )
 
+    def test_qwen3_moe(self):
+        from mlx_lm.models import qwen3_moe
+
+        args = qwen3_moe.ModelArgs(
+            model_type="qwen3_moe",
+            hidden_size=1024,
+            num_hidden_layers=4,
+            intermediate_size=2048,
+            num_attention_heads=4,
+            num_key_value_heads=4,
+            rms_norm_eps=1e-5,
+            head_dim=128,
+            vocab_size=10_000,
+            decoder_sparse_step=1,
+            mlp_only_layers=[],
+            num_experts_per_tok=4,
+            num_experts=16,
+            moe_intermediate_size=1024,
+            rope_theta=1000,
+            max_position_embeddings=4096,
+            tie_word_embeddings=False,
+            norm_topk_prob=True,
+        )
+        model = qwen3_moe.Model(args)
+        self.model_test_runner(
+            model, args.model_type, args.vocab_size, args.num_hidden_layers
+        )
+
+    def test_qwen3(self):
+        from mlx_lm.models import qwen3
+
+        args = qwen3.ModelArgs(
+            model_type="qwen3",
+            hidden_size=1024,
+            num_hidden_layers=4,
+            intermediate_size=2048,
+            num_attention_heads=4,
+            num_key_value_heads=4,
+            rms_norm_eps=1e-5,
+            vocab_size=10_000,
+            head_dim=128,
+            max_position_embeddings=4096,
+            tie_word_embeddings=False,
+            rope_theta=1000,
+        )
+        model = qwen3.Model(args)
+        self.model_test_runner(
+            model, args.model_type, args.vocab_size, args.num_hidden_layers
+        )
+
     def test_qwen2_moe(self):
         from mlx_lm.models import qwen2_moe
 

tests/test_prompt_cache.py

@@ -9,6 +9,7 @@ import mlx.core as mx
 
 from mlx_lm.generate import generate_step
 from mlx_lm.models.cache import (
+    ChunkedKVCache,
     KVCache,
     MambaCache,
     QuantizedKVCache,
@@ -95,7 +96,13 @@ class TestPromptCache(unittest.TestCase):
     def test_save_load_mixed_cache(self):
         cache_file = os.path.join(self.test_dir, "prompt_cache.safetensors")
 
-        cache = [MambaCache(), KVCache(), RotatingKVCache(8), MambaCache()]
+        cache = [
+            MambaCache(),
+            KVCache(),
+            RotatingKVCache(8),
+            MambaCache(),
+            ChunkedKVCache(256),
+        ]
         for c in cache:
             if isinstance(c, MambaCache):
                 c[0] = mx.random.uniform(shape=(4, 4, 4))

tests/test_sample_utils.py

@@ -2,7 +2,7 @@ import unittest
 
 import mlx.core as mx
 
-from mlx_lm.sample_utils import apply_min_p, apply_top_k, apply_top_p
+from mlx_lm.sample_utils import apply_min_p, apply_top_k, apply_top_p, apply_xtc
 
 
 class TestSampleUtils(unittest.TestCase):
@@ -94,6 +94,28 @@ class TestSampleUtils(unittest.TestCase):
             actual_probs.tolist(), [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0]]
         )
 
+    def test_apply_xtc(self):
+        # Test the threshold
+        probs = mx.array([[0.4, 0.3, 0.15, 0.15]])
+        new_probs = mx.softmax(apply_xtc(mx.log(probs), 1, 0.2, []), -1)
+        expected = mx.array([[0, 0.5, 0.25, 0.25]])
+        self.assertTrue(mx.allclose(new_probs, expected))
+        probs = mx.array([[0.4, 0.3, 0.15, 0.15]])
+        new_probs = mx.softmax(apply_xtc(mx.log(probs), 1, 0.1, []), -1)
+        expected = mx.array([[0, 0.0, 0.5, 0.5]])
+        self.assertTrue(mx.allclose(new_probs, expected))
+
+        # Test the special tokens
+        probs = mx.array([[0.4, 0.3, 0.15, 0.15]])
+        new_probs = mx.softmax(apply_xtc(mx.log(probs), 1, 0.1, [0]), -1)
+        expected = mx.array([[4 / 7, 0.0, 1.5 / 7, 1.5 / 7]])
+        self.assertTrue(mx.allclose(new_probs, expected))
+
+        # Test that with probability 0 the probs don't change
+        probs = mx.array([[0.4, 0.3, 0.15, 0.15]])
+        new_probs = mx.softmax(apply_xtc(mx.log(probs), 0, 0.1, [0]), -1)
+        self.assertTrue(mx.allclose(new_probs, probs))
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_server.py

@@ -12,12 +12,31 @@ from mlx_lm.utils import load
 
 
 class DummyModelProvider:
-    def __init__(self):
+    def __init__(self, with_draft=False):
         HF_MODEL_PATH = "mlx-community/Qwen1.5-0.5B-Chat-4bit"
         self.model, self.tokenizer = load(HF_MODEL_PATH)
         self.model_key = (HF_MODEL_PATH, None)
 
-    def load(self, model, adapter=None):
+        # Add draft model support
+        self.draft_model = None
+        self.draft_model_key = None
+        self.cli_args = type(
+            "obj",
+            (object,),
+            {
+                "adapter_path": None,
+                "chat_template": None,
+                "use_default_chat_template": False,
+                "trust_remote_code": False,
+            },
+        )
+
+        if with_draft:
+            # Use the same model as the draft model for testing
+            self.draft_model, _ = load(HF_MODEL_PATH)
+            self.draft_model_key = HF_MODEL_PATH
+
+    def load(self, model, adapter=None, draft_model=None):
         assert model in ["default_model", "chat_model"]
         return self.model, self.tokenizer
 
@@ -130,5 +149,324 @@ class TestServer(unittest.TestCase):
         self.assertFalse(sequence_overlap([1, 2, 3], [4, 1, 2, 3]))
 
 
+class TestServerWithDraftModel(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.model_provider = DummyModelProvider(with_draft=True)
+        cls.server_address = ("localhost", 0)
+        cls.httpd = http.server.HTTPServer(
+            cls.server_address,
+            lambda *args, **kwargs: APIHandler(cls.model_provider, *args, **kwargs),
+        )
+        cls.port = cls.httpd.server_port
+        cls.server_thread = threading.Thread(target=cls.httpd.serve_forever)
+        cls.server_thread.daemon = True
+        cls.server_thread.start()
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.httpd.shutdown()
+        cls.httpd.server_close()
+        cls.server_thread.join()
+
+    def test_handle_completions_with_draft_model(self):
+        url = f"http://localhost:{self.port}/v1/completions"
+
+        post_data = {
+            "model": "default_model",
+            "prompt": "Once upon a time",
+            "max_tokens": 10,
+            "temperature": 0.0,
+            "top_p": 1.0,
+        }
+
+        response = requests.post(url, json=post_data)
+        self.assertEqual(response.status_code, 200)
+
+        response_body = json.loads(response.text)
+        self.assertIn("id", response_body)
+        self.assertIn("choices", response_body)
+        self.assertIn("usage", response_body)
+
+        # Check that tokens were generated
+        self.assertTrue(response_body["usage"]["completion_tokens"] > 0)
+
+    def test_handle_chat_completions_with_draft_model(self):
+        url = f"http://localhost:{self.port}/v1/chat/completions"
+
+        chat_post_data = {
+            "model": "chat_model",
+            "max_tokens": 10,
+            "temperature": 0.0,
+            "messages": [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": "Hello!"},
+            ],
+        }
+
+        response = requests.post(url, json=chat_post_data)
+        self.assertEqual(response.status_code, 200)
+
+        response_body = json.loads(response.text)
+        self.assertIn("id", response_body)
+        self.assertIn("choices", response_body)
+        self.assertIn("usage", response_body)
+
+        # Check that tokens were generated
+        self.assertTrue(response_body["usage"]["completion_tokens"] > 0)
+
+    def test_streaming_with_draft_model(self):
+        url = f"http://localhost:{self.port}/v1/chat/completions"
+
+        chat_post_data = {
+            "model": "chat_model",
+            "max_tokens": 10,
+            "temperature": 0.0,
+            "stream": True,
+            "messages": [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": "Hello!"},
+            ],
+        }
+
+        response = requests.post(url, json=chat_post_data, stream=True)
+        self.assertEqual(response.status_code, 200)
+
+        chunk_count = 0
+        for chunk in response.iter_lines():
+            if chunk:
+                data = chunk.decode("utf-8")
+                if data.startswith("data: ") and data != "data: [DONE]":
+                    chunk_data = json.loads(data[6:])  # Skip the "data: " prefix
+                    self.assertIn("choices", chunk_data)
+                    self.assertEqual(len(chunk_data["choices"]), 1)
+                    self.assertIn("delta", chunk_data["choices"][0])
+                    chunk_count += 1
+
+        # Make sure we got some streaming chunks
+        self.assertGreater(chunk_count, 0)
+
+    def test_prompt_cache_with_draft_model(self):
+        url = f"http://localhost:{self.port}/v1/chat/completions"
+
+        # First request to initialize cache
+        chat_post_data = {
+            "model": "chat_model",
+            "max_tokens": 5,
+            "temperature": 0.0,
+            "messages": [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": "Tell me a story about"},
+            ],
+        }
+
+        first_response = requests.post(url, json=chat_post_data)
+        self.assertEqual(first_response.status_code, 200)
+
+        # Second request with same prefix should use cache
+        chat_post_data = {
+            "model": "chat_model",
+            "max_tokens": 5,
+            "temperature": 0.0,
+            "messages": [
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": "Tell me a story about dragons."},
+            ],
+        }
+
+        second_response = requests.post(url, json=chat_post_data)
+        self.assertEqual(second_response.status_code, 200)
+
+        # Both responses should have content
+        first_response_body = json.loads(first_response.text)
+        second_response_body = json.loads(second_response.text)
+
+        self.assertIn("choices", first_response_body)
+        self.assertIn("choices", second_response_body)
+        self.assertIn("message", first_response_body["choices"][0])
+        self.assertIn("message", second_response_body["choices"][0])
+        self.assertIn("content", first_response_body["choices"][0]["message"])
+        self.assertIn("content", second_response_body["choices"][0]["message"])
+
+        # Ensure both generated content
+        self.assertIsNotNone(first_response_body["choices"][0]["message"]["content"])
+        self.assertIsNotNone(second_response_body["choices"][0]["message"]["content"])
+
+
+# --- Tests for get_prompt_cache ---
+
+from unittest.mock import MagicMock, patch
+
+from mlx_lm.server import PromptCache
+
+
+class TestGetPromptCache(unittest.TestCase):
+
+    def setUp(self):
+        """Set up mocks and a handler instance for each test."""
+        self.mock_model_provider = MagicMock()
+        # Simulate tokenizer needed for decoding in original debug logs (though not strictly needed for cache logic)
+        self.mock_model_provider.tokenizer = MagicMock()
+        self.mock_model_provider.tokenizer.decode = lambda x: f"decoded({x})"
+        self.mock_model_provider.model_key = ("model_v1", None, None)
+        self.mock_model_provider.draft_model = None  # Start without draft model
+
+        # --- Prevent BaseHTTPRequestHandler.__init__ from running ---
+        # It tries to handle a request immediately, which fails with mocks.
+        # We only need the APIHandler instance with its attributes set.
+        with patch(
+            "http.server.BaseHTTPRequestHandler.__init__", lambda *args, **kwargs: None
+        ):
+            # APIHandler init still requires args for BaseHTTPRequestHandler signature,
+            # but they won't be used by the patched __init__.
+            mock_request = MagicMock()
+            mock_client_address = ("127.0.0.1", 8080)
+            mock_server = MagicMock()
+
+            self.prompt_cache_instance = PromptCache()
+            self.handler = APIHandler(
+                self.mock_model_provider,
+                mock_request,
+                mock_client_address,
+                mock_server,
+                prompt_cache=self.prompt_cache_instance,  # Inject our cache instance
+            )
+            # Manually set attributes usually set by APIHandler.__init__ if needed
+            # self.handler.created = MagicMock()
+            # self.handler.system_fingerprint = MagicMock()
+            # (Not strictly necessary for get_prompt_cache testing)
+
+    @patch("mlx_lm.server.make_prompt_cache")
+    def test_initial_request_empty_cache(self, mock_make_cache):
+        """Test first request when the cache is empty."""
+        mock_make_cache.return_value = "new_cache_obj"
+        prompt = [1, 2, 3]
+
+        processed_prompt = self.handler.get_prompt_cache(prompt)
+
+        self.assertEqual(processed_prompt, [1, 2, 3])
+        self.assertEqual(self.handler.prompt_cache.tokens, [1, 2, 3])
+        self.assertEqual(self.handler.prompt_cache.cache, "new_cache_obj")
+        self.assertEqual(self.handler.prompt_cache.model_key, ("model_v1", None, None))
+        mock_make_cache.assert_called_once()
+
+    def test_identical_request_full_hit(self):
+        """Test when the new prompt is identical to the cached one."""
+        self.handler.prompt_cache.tokens = [1, 2, 3]
+        self.handler.prompt_cache.model_key = ("model_v1", None, None)
+        self.handler.prompt_cache.cache = "existing_cache_obj"
+        prompt = [1, 2, 3]
+
+        # Mock common_prefix_len to return the full length
+        with patch("mlx_lm.server.common_prefix_len", return_value=3):
+            processed_prompt = self.handler.get_prompt_cache(prompt)
+
+        # Should process nothing, cache remains unchanged
+        self.assertEqual(processed_prompt, [])
+        self.assertEqual(self.handler.prompt_cache.tokens, [1, 2, 3])
+        self.assertEqual(self.handler.prompt_cache.cache, "existing_cache_obj")
+
+    def test_cache_is_prefix(self):
+        """Test when the cached prompt is a prefix of the new prompt."""
+        self.handler.prompt_cache.tokens = [1, 2, 3]
+        self.handler.prompt_cache.model_key = ("model_v1", None, None)
+        self.handler.prompt_cache.cache = "existing_cache_obj"
+        prompt = [1, 2, 3, 4, 5]
+
+        with patch("mlx_lm.server.common_prefix_len", return_value=3):
+            processed_prompt = self.handler.get_prompt_cache(prompt)
+
+        # Should process the suffix, cache tokens updated
+        self.assertEqual(processed_prompt, [4, 5])
+        self.assertEqual(self.handler.prompt_cache.tokens, [1, 2, 3, 4, 5])
+        self.assertEqual(self.handler.prompt_cache.cache, "existing_cache_obj")
+
+    @patch("mlx_lm.server.trim_prompt_cache")
+    @patch("mlx_lm.server.can_trim_prompt_cache", return_value=True)
+    def test_partial_match_trim_success(self, mock_can_trim, mock_trim_cache):
+        """Test partial match where cache is longer and trimming succeeds."""
+        self.handler.prompt_cache.tokens = [1, 2, 3, 4, 5]
+        self.handler.prompt_cache.model_key = ("model_v1", None, None)
+        self.handler.prompt_cache.cache = "existing_cache_obj"
+        prompt = [1, 2, 3, 6, 7]  # Diverges after token 3
+
+        with patch("mlx_lm.server.common_prefix_len", return_value=3):
+            processed_prompt = self.handler.get_prompt_cache(prompt)
+
+        # Should process the new suffix, cache trimmed and updated
+        self.assertEqual(processed_prompt, [6, 7])
+        self.assertEqual(self.handler.prompt_cache.tokens, [1, 2, 3, 6, 7])
+        mock_can_trim.assert_called_once_with("existing_cache_obj")
+        # Called with cache object and num_to_trim (5 - 3 = 2)
+        mock_trim_cache.assert_called_once_with("existing_cache_obj", 2)
+        self.assertEqual(
+            self.handler.prompt_cache.cache, "existing_cache_obj"
+        )  # Cache obj itself isn't changed by mock
+
+    @patch("mlx_lm.server.make_prompt_cache")
+    @patch("mlx_lm.server.trim_prompt_cache")
+    @patch("mlx_lm.server.can_trim_prompt_cache", return_value=False)
+    def test_partial_match_trim_fail(
+        self, mock_can_trim, mock_trim_cache, mock_make_cache
+    ):
+        """Test partial match where cache is longer but trimming fails."""
+        mock_make_cache.return_value = "new_cache_obj_on_reset"
+        self.handler.prompt_cache.tokens = [1, 2, 3, 4, 5]
+        self.handler.prompt_cache.model_key = ("model_v1", None, None)
+        self.handler.prompt_cache.cache = "existing_cache_obj"
+        prompt = [1, 2, 3, 6, 7]  # Diverges after token 3
+
+        with patch("mlx_lm.server.common_prefix_len", return_value=3):
+            processed_prompt = self.handler.get_prompt_cache(prompt)
+
+        # Should process the full prompt, cache reset
+        self.assertEqual(processed_prompt, [1, 2, 3, 6, 7])
+        self.assertEqual(self.handler.prompt_cache.tokens, [1, 2, 3, 6, 7])
+        mock_can_trim.assert_called_once_with("existing_cache_obj")
+        mock_trim_cache.assert_not_called()
+        mock_make_cache.assert_called_once()  # Cache was reset
+        self.assertEqual(self.handler.prompt_cache.cache, "new_cache_obj_on_reset")
+
+    @patch("mlx_lm.server.make_prompt_cache")
+    def test_no_common_prefix(self, mock_make_cache):
+        """Test when there is no common prefix between cache and prompt."""
+        mock_make_cache.return_value = "new_cache_obj"
+        self.handler.prompt_cache.tokens = [1, 2, 3]
+        self.handler.prompt_cache.model_key = ("model_v1", None, None)
+        self.handler.prompt_cache.cache = "existing_cache_obj"
+        prompt = [4, 5, 6]
+
+        with patch("mlx_lm.server.common_prefix_len", return_value=0):
+            processed_prompt = self.handler.get_prompt_cache(prompt)
+
+        # Should process the full prompt, cache reset
+        self.assertEqual(processed_prompt, [4, 5, 6])
+        self.assertEqual(self.handler.prompt_cache.tokens, [4, 5, 6])
+        mock_make_cache.assert_called_once()
+        self.assertEqual(self.handler.prompt_cache.cache, "new_cache_obj")
+
+    @patch("mlx_lm.server.make_prompt_cache")
+    def test_model_changed(self, mock_make_cache):
+        """Test cache reset when the model key changes."""
+        mock_make_cache.return_value = "new_cache_obj_model_change"
+        self.handler.prompt_cache.tokens = [1, 2, 3]
+        self.handler.prompt_cache.model_key = ("model_v1", None, None)  # Original key
+        self.handler.prompt_cache.cache = "existing_cache_obj"
+
+        # Simulate model provider having a new key
+        self.mock_model_provider.model_key = ("model_v2", None, None)
+        prompt = [1, 2, 3, 4]
+
+        # No need to mock common_prefix_len, model check happens first
+        processed_prompt = self.handler.get_prompt_cache(prompt)
+
+        # Should process the full prompt, cache reset
+        self.assertEqual(processed_prompt, [1, 2, 3, 4])
+        self.assertEqual(self.handler.prompt_cache.tokens, [1, 2, 3, 4])
+        mock_make_cache.assert_called_once()
+        self.assertEqual(self.handler.prompt_cache.cache, "new_cache_obj_model_change")
+        self.assertEqual(self.handler.prompt_cache.model_key, ("model_v2", None, None))
+
+
 if __name__ == "__main__":
     unittest.main()