Replace recursion with an explicit stack

Add DP search across multiple quant options
Fix the formating
2025-07-06 16:31:49 -07:00 · 2025-07-06 16:31:47 -07:00 · 2025-07-06 16:11:38 -07:00 · 2025-07-06 15:54:58 -07:00 · 2025-07-04 19:18:31 -07:00 · 2025-07-02 16:34:07 -07:00
81 changed files with 3825 additions and 719 deletions
@@ -8,5 +8,5 @@ with a short description of your contribution(s) below. For example:
 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`, 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`.
+- 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`, Rednote `dots.llm1`, 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`, and `reporting training metrics to WandB (Weights & Biases)`.
+- Prince Canuma: Helped add support for the following model architectures: HuggingFace's `Starcoder2`, Cohere's `Cohere (1 and 2)`, Alibaba Qwen's `Qwen (2, 3 and MoE)`, Microsoft's `Phi (3 and 3.5 MoE)`, `BitNet1.58`, Meta's `Llama (3 and 4)`, Google DeepMind's `Gemma 3`, and InterLM's `InternLM 2.5`.
@@ -1,21 +1,26 @@
 # Learned Quantization 

-To reduce the quality loss from quantization MLX LM has two options:
+To reduce the quality loss from quantization MLX LM has several options:

 - Distilled Weight Quantization (DWQ)
- Activation-aware Weight Quantization (AWQ)[^1].
+- Activation-aware Weight Quantization (AWQ)[^1]
+- Dynamic quantization

-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
+All methods use calibration data to tune parameters or hyper-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. Dynamic quantization estimates the
+sensitivity of a model's outputs to each layer and uses a higher precision for
+layers which have higher sensitivity.
+
+Dynamic quantization is the fastest to run. DWQ takes longer but typically
+yields better results. You can also cascade methods. For example a dynamically
+quantized model can be further refined with DWQ.

 To get started, first install the requirements:

 ```
-pip install mlx-lm[lwq]
+pip install mlx-lm[quant]
 ```

 ### DWQ
@@ -40,6 +45,57 @@ For a full list of options run:
 mlx_lm.dwq --help
 ```

+#### Tips
+
+- DWQ works best distilling to lower precision, anywhere from 2-bit to 4-bit
+  models.
+- Distilling 16-bit precision to 8-bit and even 6-bit often doesn't work well.
+  The loss starts out so low that it's difficult to reduce further.
+- Decreasing the quantization group size (e.g. `--group-size 32`) doubles the
+  number of tunable parameters and can work much better.
+- If the loss is oscillating and not going down consistently, try reducing the
+  learning rate. If it is decreasing but slowly, try increasing the learning
+  rate.
+- As a rule of thumb, lower precision can benefit from a higher learning rate
+  since the loss starts out higher. Conversely, higher precision needs a lower
+  learning rate.
+
+
+#### Memory Use
+
+A few options to reduce memory use for DWQ:
+
+- Distill from an 8-bit model instead of a 16-bit model. The 8-bit
+  models are usually as good as 16-bit precision models.
+- Use a shorter maximum sequence length. The default is 2048. Using
+  `--max-seq-length 512` reduces the memory and still gets good results.
+- Use a smaller batch size, e.g. `--batch-size 1`
+
+### Dynamic Quantization
+
+Use `mlx_lm.dynamic_quant` to generate a dynamic quantization of given model.
+For example:
+
+```bash
+mlx_lm.dynamic_quant --model mistralai/Mistral-7B-Instruct-v0.3
+```
+
+The script will estimate the sensitivity for each quantizable layer in the
+model. It will then quantize the model using higher precision (default 5 bits)
+for the more sensitive layers and lower precision (default 4 bits) for the
+rest. The script also saves a JSON file with each layer's sensitivities which
+saves needing to compute it multiple times to make different precision quants
+of the same model.
+
+Some important options are:
+
+- `--target-bpw`: The target bits-per-weight. For a given set of quantization
+  parameters only certain ranges are possible. For example, with the default
+  parameters a BPW in the range `[4.5, 5.5]` is achievable.
+- `--sensitivities`: A path to a precomputed sensitivities file.
+- `--low-bits`: The number of bits to use for the less sensitive layers.
+- `--high-bits`: The number of bits to use for the more sensitive layers.
+
 ### AWQ 

 Use `mlx_lm.awq` to run AWQ on a given model. For example:
@@ -76,6 +76,11 @@ You can specify the output location with `--adapter-path`.
 You can resume fine-tuning with an existing adapter with
 `--resume-adapter-file <path_to_adapters.safetensors>`.

+#### Logging
+
+You can log training metrics to Weights & Biases by passing a project name with
+the `--wandb` flag. Make sure to install wandb with `pip install wandb`.
+
 #### Prompt Masking

 The default training computes a loss for every token in the sample. You can
@@ -379,7 +384,7 @@ mlx_lm.lora \
    --train \
    --batch-size 1 \
    --num-layers 4 \
-    --data wikisql
+    --data  mlx-community/wikisql
 ```

 The above command on an M1 Max with 32 GB runs at about 250
@@ -1,50 +0,0 @@
-# Model Merging
-
-You can use `mlx-lm` to merge models and upload them to the Hugging
-Face hub or save them locally for LoRA fine tuning.
-
-The main command is `mlx_lm.merge`:
-
-```shell
-mlx_lm.merge --config config.yaml 
-```
-
-The merged model will be saved by default in `mlx_merged_model`. To see a
-full list of options run:
-
-```shell
-mlx_lm.merge --help
-```
-
-Here is an example `config.yaml`:
-
-```yaml
-models:
-  - OpenPipe/mistral-ft-optimized-1218
-  - mlabonne/NeuralHermes-2.5-Mistral-7B
-method: slerp
-parameters:
-  t:
-    - filter: self_attn
-      value: [0, 0.5, 0.3, 0.7, 1]
-    - filter: mlp
-      value: [1, 0.5, 0.7, 0.3, 0]
-    - value: 0.5
-```
-
-The `models` field is a list of Hugging Face repo ids. The first model in the
-list is treated as the base model into which the remaining models are merged.
-
-The `method` field is the merging method. Right now `slerp` is the only
-supported method.
-
-The `parameters` are the corresponding parameters for the given `method`.
-Each parameter is a list with `filter` determining which layer the parameter
-applies to and `value` determining the actual value used. The last item in
-the list without a `filter` field is the default.
-
-If `value` is a list, it specifies the start and end values for the
-corresponding segment of blocks. In the example above, the models have 32
-blocks. For blocks 1-8, the layers with `self_attn` in the name will use the
-values `np.linspace(0, 0.5, 8)`, the same layers in the next 8 blocks (9-16)
-will use `np.linspace(0.5, 0.3, 8)`, and so on.
@@ -54,18 +54,24 @@ curl localhost:8080/v1/chat/completions \
  sequences of tokens on which the generation should stop.

 - `max_tokens`: (Optional) An integer specifying the maximum number of tokens
-  to generate. Defaults to `100`.
+  to generate. Defaults to `512`.

 - `stream`: (Optional) A boolean indicating if the response should be
  streamed. If true, responses are sent as they are generated. Defaults to
  false.

 - `temperature`: (Optional) A float specifying the sampling temperature.
-  Defaults to `1.0`.
+  Defaults to `0.0`.

 - `top_p`: (Optional) A float specifying the nucleus sampling parameter.
  Defaults to `1.0`.

+- `top_k`: (Optional) An integer specifying the top-k sampling parameter.
+  Defaults to `0` (disabled).
+
+- `min_p`: (Optional) A float specifying the min-p sampling parameter.
+  Defaults to `0.0` (disabled).
+
 - `repetition_penalty`: (Optional) Applies a penalty to repeated tokens.
  Defaults to `1.0`.

@@ -92,7 +98,6 @@ curl localhost:8080/v1/chat/completions \
 - `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.
@@ -5,8 +5,9 @@ import sys

 if __name__ == "__main__":
    subcommands = {
-        "awq",
-        "dwq",
+        "quant.awq",
+        "quant.dwq",
+        "quant.dynamic_quant",
        "cache_prompt",
        "chat",
        "convert",
@@ -14,7 +15,6 @@ if __name__ == "__main__":
        "fuse",
        "generate",
        "lora",
-        "merge",
        "server",
        "manage",
        "upload",
@@ -1,3 +1,3 @@
 # Copyright © 2023-2024 Apple Inc.

-__version__ = "0.24.0"
+__version__ = "0.25.3"
@@ -1,7 +1,6 @@
 # Copyright © 2023-2024 Apple Inc.

 import argparse
-import json

 import mlx.core as mx

@@ -6,7 +6,7 @@ from typing import Callable, Optional, Union

 import mlx.core as mx
 import mlx.nn as nn
-from mlx.utils import tree_flatten
+from mlx.utils import tree_map_with_path

 from .utils import (
    dequantize_model,
@@ -59,6 +59,8 @@ def mixed_quant_predicate_builder(

        if not hasattr(module, "to_quantized"):
            return False
+        if module.weight.shape[1] % group_size != 0:
+            return False

        index = (
            int(path.split(".")[layer_location])
@@ -112,47 +114,54 @@ def convert(
        )

    print("[INFO] Loading")
-    model_path = get_model_path(hf_path, revision=revision)
+    model_path, hf_path = get_model_path(hf_path, revision=revision)
    model, config, tokenizer = fetch_from_hub(model_path, lazy=True)

+    def base_quant_predicate(path, module, config):
+        if not hasattr(module, "to_quantized"):
+            return False
+        if module.weight.shape[1] % q_group_size != 0:
+            return False
+        return True
+
    if isinstance(quant_predicate, str):
        quant_predicate = mixed_quant_predicate_builder(quant_predicate, model)
+    quant_predicate = quant_predicate or base_quant_predicate

    if dtype is None:
        dtype = config.get("torch_dtype", None)
-    weights = dict(tree_flatten(model.parameters()))
    if dtype in MODEL_CONVERSION_DTYPES:
        print("[INFO] Using dtype:", dtype)
        dtype = getattr(mx, dtype)
+        cast_predicate = getattr(model, "cast_predicate", lambda _: True)

-        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()
-        }
+        def set_dtype(k, v):
+            if cast_predicate(k) and mx.issubdtype(v.dtype, mx.floating):
+                return v.astype(dtype)
+            else:
+                return v
+
+        model.update(tree_map_with_path(set_dtype, model.parameters()))

    if quantize and dequantize:
        raise ValueError("Choose either quantize or dequantize, not both.")

    if quantize:
        print("[INFO] Quantizing")
-        model.load_weights(list(weights.items()))
-        weights, config = quantize_model(
+        model, config = quantize_model(
            model, config, q_group_size, q_bits, quant_predicate=quant_predicate
        )

    if dequantize:
        print("[INFO] Dequantizing")
+        config.pop("quantization", None)
+        config.pop("quantization_config", None)
        model = dequantize_model(model)
-        weights = dict(tree_flatten(model.parameters()))

-    del model
    save(
        mlx_path,
        model_path,
-        weights,
+        model,
        tokenizer,
        config,
        hf_repo=hf_path,
@@ -1,5 +1,5 @@
 # The path to the local model directory or Hugging Face repo.
-model: "mlx_model"
+model: "mlx-community/Llama-3.2-1B-Instruct"

 # Whether or not to train (boolean)
 train: true
@@ -17,7 +17,7 @@ optimizer: adamw
 #     bias_correction: true

 # Directory with {train, valid, test}.jsonl files
-data: "/path/to/training/data"
+data: "mlx-community/WikiSQL"

 # The PRNG seed
 seed: 0
@@ -37,6 +37,9 @@ val_batches: 25
 # Adam learning rate.
 learning_rate: 1e-5

+# Whether to report the logs to WandB
+# wand: "wandb-project"
+
 # Number of training steps between loss reporting.
 steps_per_report: 10

@@ -0,0 +1,65 @@
+# Copyright © 2025 Apple Inc.
+"""
+This is an example of tool use with mlx_lm and the OpenAI client.
+
+To run, first start the server:
+
+>>> mlx_lm.server
+
+Then run this script.
+"""
+from openai import OpenAI
+
+client = OpenAI(base_url="http://localhost:8080/v1", api_key="not-needed")
+
+model = "mlx-community/qwen3-4b-4bit-DWQ"
+messages = [{"role": "user", "content": "What's the weather in Boston?"}]
+
+tools = [
+    {
+        "type": "function",
+        "function": {
+            "name": "get_current_weather",
+            "description": "Get the current weather in a given location",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "location": {
+                        "type": "string",
+                        "description": "The city and state, e.g. San Francisco, CA",
+                    },
+                    "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
+                },
+                "required": ["location"],
+            },
+        },
+    }
+]
+
+
+def get_current_weather(**kwargs):
+    return "51 Farenheit, clear skies"
+
+
+functions = {"get_current_weather": get_current_weather}
+
+# The first query generates a tool call
+response = client.chat.completions.create(
+    model=model,
+    messages=messages,
+    tools=tools,
+)
+
+# Call the function
+function = response.choices[0].message.tool_calls[0].function
+tool_result = functions[function.name](**function.arguments)
+
+# Put the result of the function in the messages and generate the final
+# response:
+messages.append({"role": "tool", "name": function.name, "content": tool_result})
+response = client.chat.completions.create(
+    model=model,
+    messages=messages,
+    tools=tools,
+)
+print(response.choices[0].message.content)
@@ -4,8 +4,6 @@ from pathlib import Path
 from mlx.utils import tree_flatten, tree_unflatten

 from .gguf import convert_to_gguf
-from .tuner.dora import DoRAEmbedding, DoRALinear
-from .tuner.lora import LoRAEmbedding, LoRALinear, LoRASwitchLinear
 from .tuner.utils import dequantize, load_adapters
 from .utils import (
    fetch_from_hub,
@@ -35,12 +33,6 @@ def parse_arguments() -> argparse.Namespace:
        default="adapters",
        help="Path to the trained adapter weights and config.",
    )
-    parser.add_argument(
-        "--hf-path",
-        type=str,
-        default=None,
-        help="Path to the original Hugging Face model. Required for upload if --model is a local directory.",
-    )
    parser.add_argument(
        "--upload-repo",
        help="The Hugging Face repo to upload the model to.",
@@ -70,14 +62,16 @@ def main() -> None:
    print("Loading pretrained model")
    args = parse_arguments()

-    model_path = get_model_path(args.model)
+    model_path, hf_path = get_model_path(args.model)
    model, config, tokenizer = fetch_from_hub(model_path)

    model.freeze()
    model = load_adapters(model, args.adapter_path)

    fused_linears = [
-        (n, m.fuse()) for n, m in model.named_modules() if hasattr(m, "fuse")
+        (n, m.fuse(de_quantize=args.de_quantize))
+        for n, m in model.named_modules()
+        if hasattr(m, "fuse")
    ]

    if fused_linears:
@@ -88,18 +82,15 @@ def main() -> None:
        model = dequantize(model)
        config.pop("quantization", None)

-    weights = dict(tree_flatten(model.parameters()))
-
    save_path = Path(args.save_path)
-    hf_path = args.hf_path or (args.model if not Path(args.model).exists() else None)
    save(
        save_path,
        model_path,
-        weights,
+        model,
        tokenizer,
        config,
        hf_repo=hf_path,
-        donate_weights=False,
+        donate_model=False,
    )

    if args.export_gguf:
@@ -108,6 +99,7 @@ def main() -> None:
            raise ValueError(
                f"Model type {model_type} not supported for GGUF conversion."
            )
+        weights = dict(tree_flatten(model.parameters()))
        convert_to_gguf(model_path, weights, config, str(save_path / args.gguf_path))

    if args.upload_repo is not None:
@@ -115,7 +107,7 @@ def main() -> None:
            raise ValueError(
                "Must provide original Hugging Face repo to upload local model."
            )
-        upload_to_hub(args.save_path, args.upload_repo, hf_path)
+        upload_to_hub(args.save_path, args.upload_repo)


 if __name__ == "__main__":
@@ -26,12 +26,10 @@ from .models import cache
 from .models.cache import (
    QuantizedKVCache,
    load_prompt_cache,
-    make_prompt_cache,
-    trim_prompt_cache,
 )
 from .sample_utils import make_sampler
 from .tokenizer_utils import TokenizerWrapper
-from .utils import load
+from .utils import does_model_support_input_embeddings, load

 DEFAULT_PROMPT = "hello"
 DEFAULT_MAX_TOKENS = 100
@@ -216,6 +214,12 @@ def wired_limit(model: nn.Module, streams: Optional[List[mx.Stream]] = None):
    async eval could be running pass in the streams to synchronize with prior
    to exiting the context manager.
    """
+    if not mx.metal.is_available():
+        try:
+            yield
+        finally:
+            return
+
    model_bytes = tree_reduce(
        lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, model, 0
    )
@@ -231,7 +235,7 @@ def wired_limit(model: nn.Module, streams: Optional[List[mx.Stream]] = None):
        )
    old_limit = mx.set_wired_limit(max_rec_size)
    try:
-        yield None
+        yield
    finally:
        if streams is not None:
            for s in streams:
@@ -298,6 +302,7 @@ def generate_step(
    kv_group_size: int = 64,
    quantized_kv_start: int = 0,
    prompt_progress_callback: Optional[Callable[int, int]] = None,
+    input_embeddings: Optional[mx.array] = None,
 ) -> Generator[Tuple[mx.array, mx.array], None, None]:
    """
    A generator producing token ids based on the given prompt from the model.
@@ -322,14 +327,22 @@ def generate_step(
        kv_group_size (int): Group size for KV cache quantization. Default: ``64``.
        quantized_kv_start (int): Step to begin using a quantized KV cache.
           when ``kv_bits`` is non-None. Default: ``0``.
-        prompt_prorgress_callback (Callable[int, int]): A call-back which takes the
+        prompt_progress_callback (Callable[int, int]): A call-back which takes the
           prompt tokens processed so far and the total number of prompt tokens.
+        input_embeddings (mx.array, optional): Input embeddings to use in place of
+           prompt tokens. Default: ``None``.

    Yields:
        Tuple[mx.array, mx.array]: One token and a vector of log probabilities.
    """
+    if input_embeddings is not None:
+        if not does_model_support_input_embeddings(model):
+            raise ValueError("Model does not support input embeddings.")
+        if len(prompt) != 0:
+            raise ValueError(
+                "If using input embeddings, prompt tokens must be an empty array."
+            )

-    y = prompt
    tokens = None

    # Create the KV cache for generation
@@ -338,8 +351,6 @@ def generate_step(
            model,
            max_kv_size=max_kv_size,
        )
-    elif len(prompt_cache) != len(model.layers):
-        raise ValueError("Wrong number of layers in the prompt cache.")

    prompt_progress_callback = prompt_progress_callback or (lambda *_: None)

@@ -352,15 +363,22 @@ def generate_step(

    sampler = sampler or (lambda x: mx.argmax(x, axis=-1))

+    def _model_call(y):
+        if y.ndim == 3:
+            return model(None, cache=prompt_cache, input_embeddings=y)
+        else:
+            return model(y, cache=prompt_cache)
+
    def _step(y):
+        nonlocal tokens
+
        with mx.stream(generation_stream):
-            logits = model(y[None], cache=prompt_cache)
+            logits = _model_call(y[None])
+
            logits = logits[:, -1, :]

-            if logits_processors:
-                nonlocal tokens
+            if logits_processors and input_embeddings is None:
                tokens = mx.concat([tokens, y]) if tokens is not None else y
-
                for processor in logits_processors:
                    logits = processor(tokens, logits)

@@ -370,11 +388,14 @@ def generate_step(
            y = sampler(logprobs)
            return y, logprobs.squeeze(0)

+    using_embeddings = input_embeddings is not None
+
+    y = input_embeddings if using_embeddings else prompt
    with mx.stream(generation_stream):
-        total_prompt_tokens = y.size
+        total_prompt_tokens = y.shape[0]
        prompt_processed_tokens = 0
-        while y.size > prefill_step_size:
-            model(y[:prefill_step_size][None], cache=prompt_cache)
+        while y.shape[0] > prefill_step_size:
+            _model_call(y[:prefill_step_size][None])
            quantize_cache_fn(prompt_cache)
            mx.eval([c.state for c in prompt_cache])
            prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens)
@@ -454,8 +475,6 @@ def speculative_generate_step(
    if prompt_cache is None:
        model_cache = cache.make_prompt_cache(model)
        draft_cache = cache.make_prompt_cache(draft_model)
-    elif len(prompt_cache) != (len(model.layers) + len(draft_model.layers)):
-        raise ValueError("Wrong number of layers in the prompt cache.")
    else:
        model_cache = prompt_cache[: len(model.layers)]
        draft_cache = prompt_cache[len(model.layers) :]
@@ -1,5 +1,3 @@
-# Copyright © 2024 Apple Inc.
-
 import argparse
 import math
 import os
@@ -13,6 +11,7 @@ import mlx.optimizers as optim
 import numpy as np
 import yaml

+from .tuner.callbacks import WandBCallback
 from .tuner.datasets import CacheDataset, load_dataset
 from .tuner.trainer import TrainingArgs, TrainingCallback, evaluate, train
 from .tuner.utils import (
@@ -66,8 +65,9 @@ CONFIG_DEFAULTS = {
    "config": None,
    "grad_checkpoint": False,
    "lr_schedule": None,
-    "lora_parameters": {"rank": 8, "dropout": 0.0, "scale": 10.0},
+    "lora_parameters": {"rank": 8, "dropout": 0.0, "scale": 20.0},
    "mask_prompt": False,
+    "wandb": None,
 }


@@ -179,6 +179,12 @@ def build_parser():
        help="Use gradient checkpointing to reduce memory use.",
        default=None,
    )
+    parser.add_argument(
+        "--wandb",
+        type=str,
+        default=None,
+        help="WandB project name to report training metrics. Disabled if None.",
+    )
    parser.add_argument("--seed", type=int, help="The PRNG seed")
    return parser

@@ -201,6 +207,8 @@ def train_model(
    if args.fine_tune_type == "full":
        for l in model.layers[-max(args.num_layers, 0) :]:
            l.unfreeze()
+
+        args.lora_parameters = None
    elif args.fine_tune_type in ["lora", "dora"]:
        # Convert linear layers to lora/dora layers and unfreeze in the process
        linear_to_lora_layers(
@@ -281,6 +289,14 @@ def evaluate_model(args, model: nn.Module, test_set):
 def run(args, training_callback: TrainingCallback = None):
    np.random.seed(args.seed)

+    if args.wandb is not None:
+        training_callback = WandBCallback(
+            project_name=args.wandb,
+            log_dir=args.adapter_path,
+            config=vars(args),
+            wrapped_callback=training_callback,
+        )
+
    print("Loading pretrained model")
    model, tokenizer = load(args.model)

@@ -1,176 +0,0 @@
-# Copyright © 2023-2024 Apple Inc.
-
-import argparse
-import glob
-import shutil
-from pathlib import Path
-from typing import Optional
-
-import mlx.core as mx
-import mlx.nn as nn
-import numpy as np
-import yaml
-from mlx.utils import tree_flatten, tree_map
-
-from .utils import (
-    fetch_from_hub,
-    get_model_path,
-    save_config,
-    save_weights,
-    upload_to_hub,
-)
-
-
-def configure_parser() -> argparse.ArgumentParser:
-    """
-    Configures and returns the argument parser for the script.
-
-    Returns:
-        argparse.ArgumentParser: Configured argument parser.
-    """
-    parser = argparse.ArgumentParser(description="Merge multiple models.")
-
-    parser.add_argument("--config", type=str, help="Path to the YAML config.")
-    parser.add_argument(
-        "--mlx-path",
-        type=str,
-        default="mlx_merged_model",
-        help="Path to save the MLX model.",
-    )
-    parser.add_argument(
-        "--upload-repo",
-        help="The Hugging Face repo to upload the model to.",
-        type=str,
-        default=None,
-    )
-    return parser
-
-
-def slerp(t, w1, w2, eps=1e-5):
-    """
-    Spherical linear interpolation
-
-    Args:
-        t (float): Interpolation weight in [0.0, 1.0]
-        w1 (mx.array): First input
-        w2 (mx.array): Second input
-        eps (float): Constant for numerical stability
-    Returns:
-        mx.array: Interpolated result
-    """
-    t = float(t)
-    if t == 0:
-        return w1
-    elif t == 1:
-        return w2
-    # Normalize
-    v1 = w1 / mx.linalg.norm(w1)
-    v2 = w2 / mx.linalg.norm(w2)
-    # Angle
-    dot = mx.clip((v1 * v2).sum(), 0.0, 1.0)
-    theta = mx.arccos(dot)
-    sin_theta = mx.sin(theta + eps)
-    s1 = mx.sin(theta * (1 - t)) / sin_theta
-    s2 = mx.sin(theta * t) / sin_theta
-    return s1 * w1 + s2 * w2
-
-
-def merge_models(base_model: nn.Module, model: nn.Module, config: dict):
-    method = config.get("method", None)
-    if method != "slerp":
-        raise ValueError(f"Merge method {method} not supported")
-
-    num_layers = len(model.layers)
-
-    def unpack_values(vals):
-        if isinstance(vals, (int, float)):
-            return np.full(num_layers, vals)
-        bins = len(vals) - 1
-        sizes = [num_layers // bins] * bins
-        sizes[-1] = num_layers - sum(sizes[:-1])
-        return np.concatenate(
-            [np.linspace(v1, v2, s) for v1, v2, s in zip(vals[:-1], vals[1:], sizes)]
-        )
-
-    param_list = config["parameters"]["t"]
-    params = {}
-    filter_keys = set()
-    for pl in param_list[:-1]:
-        params[pl["filter"]] = unpack_values(pl["value"])
-        filter_keys.add(pl["filter"])
-    default = unpack_values(param_list[-1]["value"])
-
-    for e in range(num_layers):
-        bl = base_model.layers[e]
-        l = model.layers[e]
-        base_weights = bl.parameters()
-        weights = l.parameters()
-        for k, w1 in base_weights.items():
-            w2 = weights[k]
-            t = params.get(k, default)[e]
-            base_weights[k] = tree_map(lambda x, y: slerp(t, x, y), w1, w2)
-        base_model.update(base_weights)
-
-
-def merge(
-    config: str,
-    mlx_path: str = "mlx_model",
-    upload_repo: Optional[str] = None,
-):
-    with open(config, "r") as fid:
-        merge_conf = yaml.safe_load(fid)
-    print("[INFO] Loading")
-
-    model_paths = merge_conf.get("models", [])
-    if len(model_paths) < 2:
-        raise ValueError(f"Expected at least 2 models, got {len(model_paths)}.")
-
-    # Load all models
-    base_hf_path = model_paths[0]
-    base_path = get_model_path(base_hf_path)
-    base_model, base_config, tokenizer = fetch_from_hub(base_path, lazy=True)
-    models = []
-    for mp in model_paths[1:]:
-        model, model_config, _ = fetch_from_hub(get_model_path(mp), lazy=True)
-        base_type = base_config["model_type"]
-        model_type = model_config["model_type"]
-        if base_type != model_type:
-            raise ValueError(
-                f"Can only merge models of the same type,"
-                f" but got {base_type} and {model_type}."
-            )
-        models.append(model)
-
-    # Merge models into base model
-    for m in models:
-        merge_models(base_model, m, merge_conf)
-
-    # Save base model
-    mlx_path = Path(mlx_path)
-    weights = dict(tree_flatten(base_model.parameters()))
-    del models, base_model
-    save_weights(mlx_path, weights, donate_weights=True)
-    py_files = glob.glob(str(base_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")
-
-    if upload_repo is not None:
-        upload_to_hub(mlx_path, upload_repo, base_hf_path)
-
-
-def main():
-    parser = configure_parser()
-    args = parser.parse_args()
-    merge(**vars(args))
-
-
-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()
@@ -0,0 +1,397 @@
+# Copyright © 2025 Apple Inc.
+
+import math
+from dataclasses import dataclass
+from functools import partial
+from itertools import accumulate
+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 ConcatenateKVCache, KVCache
+from .rope_utils import initialize_rope
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_dim: int
+    num_layers: int
+    num_kv_reuse_layers: int
+    num_heads: int
+    num_kv_heads: int
+    hidden_dim_scale_factor: float = 3.25
+    rope_theta: float = 50000
+    rms_norm_eps: float = 1e-5
+
+
+class FusedLoRALinear(nn.Module):
+    def __init__(
+        self,
+        input_dims: int,
+        output_dims: list[int],
+        r: int = 8,
+        dropout: float = 0.0,
+        scale: float = 20.0,
+    ):
+        super().__init__()
+
+        self.linear = FusedLinear(input_dims, output_dims)
+        self.dropout = nn.Dropout(p=dropout)
+        self.scale = scale
+
+        scale = 1 / math.sqrt(input_dims)
+        self.lora_a = [
+            mx.random.uniform(low=-scale, high=scale, shape=(input_dims, r))
+            for _ in output_dims
+        ]
+        self.lora_b = [mx.zeros((r, od)) for od in output_dims]
+
+    def fuse(self, de_quantize: bool = False):
+        linear = self.linear
+        weight = linear.weight
+        is_quantized = isinstance(linear, FusedQuantizedLinear)
+
+        # Use the same type as the linear weight if not quantized
+        dtype = weight.dtype
+
+        if is_quantized:
+            dtype = linear.scales.dtype
+            weight = mx.dequantize(
+                weight,
+                linear.scales,
+                linear.biases,
+                linear.group_size,
+                linear.bits,
+            )
+
+        input_dims = weight.shape[-1]
+        output_dims = linear.output_dims
+        fused_linear = FusedLinear(input_dims, output_dims)
+        fused_linear.weight = weight
+        deltas = [
+            ((self.scale * b.T) @ a.T).astype(dtype)
+            for a, b in zip(self.lora_a, self.lora_b)
+        ]
+        delta = mx.concatenate(deltas, axis=0)
+        fused_linear.weight = weight + delta
+
+        if is_quantized and not de_quantize:
+            fused_linear = fused_linear.to_quantized(linear.group_size, linear.bits)
+
+        return fused_linear
+
+    def __call__(self, x):
+        dt = x.dtype
+        y = self.linear(x)
+        x = self.dropout(x)
+        z = [(x @ a) @ b for a, b in zip(self.lora_a, self.lora_b)]
+        return tuple(yi + (self.scale * zi).astype(dt) for yi, zi in zip(y, z))
+
+
+class FusedQuantizedLinear(nn.QuantizedLinear):
+    def __init__(self, input_dims, output_dims, group_size: int = 64, bits: int = 4):
+        *indices, output_dims = accumulate(output_dims)
+        self.indices = indices
+        super().__init__(
+            input_dims, output_dims, bias=False, group_size=group_size, bits=bits
+        )
+
+    @property
+    def input_dims(self):
+        return self.scales.shape[-1] * self.group_size
+
+    @property
+    def output_dims(self):
+        indices = [0] + self.indices + [self.weight.shape[0]]
+        return [indices[i] - indices[i - 1] for i in range(1, len(indices))]
+
+    def __call__(self, x):
+        x = super().__call__(x)
+        return x.split(self.indices, axis=-1)
+
+    def to_lora(self, r: int = 8, dropout: float = 0.0, scale: float = 20.0):
+        lora_lin = FusedLoRALinear(self.input_dims, self.output_dims, r, dropout, scale)
+        lora_lin.linear = self
+        return lora_lin
+
+
+class FusedLinear(nn.Linear):
+    def __init__(self, input_dims, output_dims):
+        *indices, output_dims = accumulate(output_dims)
+        self.indices = indices
+        super().__init__(input_dims, output_dims, bias=False)
+
+    @property
+    def input_dims(self):
+        return self.weight.shape[-1]
+
+    @property
+    def output_dims(self):
+        indices = [0] + self.indices + [self.weight.shape[0]]
+        return [indices[i] - indices[i - 1] for i in range(1, len(indices))]
+
+    def __call__(self, x):
+        x = super().__call__(x)
+        return x.split(self.indices, axis=-1)
+
+    def to_quantized(self, group_size: int = 64, bits: int = 4):
+        input_dims = self.input_dims
+        output_dims = self.output_dims
+        ql = FusedQuantizedLinear(input_dims, output_dims, group_size, bits)
+        ql.weight, ql.scales, ql.biases = mx.quantize(self.weight, group_size, bits)
+
+        return ql
+
+    def to_lora(self, r: int = 8, dropout: float = 0.0, scale: float = 20.0):
+        lora_lin = FusedLoRALinear(self.input_dims, self.output_dims, r, dropout, scale)
+        lora_lin.linear = self
+        return lora_lin
+
+
+@partial(mx.compile, shapeless=True)
+def fake_8bit_quant(x, scale):
+    dt = x.dtype
+    x = x.astype(mx.float32)
+    x = (x / scale).round()
+    x = mx.clip(x, -128, 127)
+    return (x * scale).astype(dt)
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_dim
+        self.n_heads = n_heads = args.num_heads
+        self.n_kv_heads = n_kv_heads = args.num_kv_heads
+        self.head_dim = head_dim = args.hidden_dim // n_heads
+        self.scale = head_dim**-0.5
+
+        qkv_dim = (n_heads + 2 * n_kv_heads) * head_dim
+        self.qkv_proj = FusedLinear(
+            dim, [n_heads * head_dim] + 2 * [n_kv_heads * head_dim]
+        )
+        self.out_proj = nn.Linear(dim, dim, bias=False)
+        self.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            True,
+        )
+        self.q_norm = nn.RMSNorm(head_dim)
+        self.k_norm = nn.RMSNorm(head_dim)
+        self.quant_key_scale = mx.array(1.0)
+        self.quant_value_scale = mx.array(1.0)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        # Get the queries, keys and values
+        queries, keys, values = self.qkv_proj(x)
+
+        # Prepare the queries, keys and values for the attention computation
+        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.q_norm(self.rope(queries, offset=cache.offset))
+            keys = self.k_norm(self.rope(keys, offset=cache.offset))
+            keys = fake_8bit_quant(keys, self.quant_key_scale)
+            values = fake_8bit_quant(values, self.quant_value_scale)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.q_norm(self.rope(queries))
+            keys = self.k_norm(self.rope(keys))
+            keys = fake_8bit_quant(keys, self.quant_key_scale)
+            values = fake_8bit_quant(values, self.quant_value_scale)
+
+        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.out_proj(output)
+
+
+class KVReuseAttention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_dim
+        self.n_heads = n_heads = args.num_heads
+        self.head_dim = head_dim = args.hidden_dim // n_heads
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.out_proj = nn.Linear(dim, dim, bias=False)
+        self.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            True,
+        )
+        self.q_norm = nn.RMSNorm(head_dim)
+
+    def __call__(
+        self,
+        x: mx.array,
+        keys: mx.array,
+        values: mx.array,
+        mask: Optional[mx.array] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+        _, _, S, _ = keys.shape
+
+        queries = self.q_proj(x)
+        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        queries = self.q_norm(self.rope(queries, offset=S - L))
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=None, scale=self.scale, mask=mask
+        )
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.out_proj(output)
+
+
+@partial(mx.compile, shapeless=True)
+def _swiglu(g, x):
+    return nn.silu(g) * x
+
+
+class MLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_dim
+        hidden_dim = int(dim * args.hidden_dim_scale_factor)
+
+        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:
+        g = self.gate_proj(x)
+        x = self.up_proj(x)
+        return self.down_proj(_swiglu(g, x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.self_attn = Attention(args)
+        self.mlp = MLP(args)
+        self.input_layernorm = nn.RMSNorm(args.hidden_dim, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_dim, eps=args.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        out = h + r
+        return out
+
+
+class KVReuseTransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.self_attn = KVReuseAttention(args)
+        self.mlp = MLP(args)
+        self.input_layernorm = nn.RMSNorm(args.hidden_dim, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_dim, eps=args.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        keys: mx.array,
+        values: mx.array,
+        mask: Optional[mx.array] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), keys, values, mask)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        out = h + r
+        return out
+
+
+class AFMModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+
+        self.embedding = nn.Embedding(args.vocab_size, args.hidden_dim)
+        self.layers = [
+            TransformerBlock(args)
+            for _ in range(args.num_layers - args.num_kv_reuse_layers)
+        ]
+        self.kv_reuse_layers = [
+            KVReuseTransformerBlock(args) for _ in range(args.num_kv_reuse_layers)
+        ]
+        self.output_norm = nn.RMSNorm(args.hidden_dim, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        h = self.embedding(inputs)
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+            cache[-1] = ConcatenateKVCache()
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, cache=c)
+
+        keys, values = cache[-1].state
+        for layer in self.kv_reuse_layers:
+            h = layer(h, keys, values, mask)
+
+        return self.output_norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = AFMModel(args)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        out = self.model(inputs, mask, cache)
+        out = self.model.embedding.as_linear(out)
+        return out
+
+    def make_cache(self):
+        return [KVCache() for _ in range(len(self.model.layers))]
+
+    @property
+    def layers(self):
+        return self.model.layers + self.model.kv_reuse_layers
@@ -0,0 +1,131 @@
+# Copyright © 2025 Apple Inc.
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.nn.layers.quantized import QuantizedLinear
+
+
+def make_bitlinear_kernel():
+    """
+    Custom Metal kernel that performs matrix multiplication directly on
+    packed weights and scales the output. This eliminates the need to
+    store unpacked weights in memory.
+    """
+    source = """
+    constexpr int M = 4;
+    constexpr int BLOCK = 32;
+
+    uint tid = thread_position_in_grid.y;
+    uint in_offset = thread_position_in_grid.x;
+
+    uint batch_idx = tid / (out_features / 4);
+    uint row_idx = tid % (out_features / 4);
+
+    float sum[4] = {0.0};
+
+    for (uint i = in_offset * M; i < in_features; i += BLOCK * M) {
+        float v[M];
+        for (int j=0; j<M; j++) {
+            v[j] = x[batch_idx * in_features + i + j];
+        }
+
+        for (int j=0; j<M; j++) {
+            uint8_t w = packed_weights[row_idx * in_features + i + j];
+            sum[0] += v[j] * ((w & 3) - 1);
+            sum[1] += v[j] * (((w >> 2) & 3) - 1);
+            sum[2] += v[j] * (((w >> 4) & 3) - 1);
+            sum[3] += v[j] * (((w >> 6) & 3) - 1);
+        }
+    }
+
+    for (int j=0; j<4; j++) {
+        sum[j] = simd_sum(sum[j]);
+    }
+
+    // Apply weight scaling by diving them or multiplying them
+    if (in_offset == 0) {
+        float scale = invert_weight_scales ? 1 / weight_scale[0] : weight_scale[0];
+        for (int i=0; i<4; i++) {
+            out[batch_idx * out_features + row_idx + i * (out_features/4)] = static_cast<T>(sum[i] * scale);
+        }
+    }
+    """
+
+    return mx.fast.metal_kernel(
+        name="bitlinear_matmul",
+        input_names=["x", "packed_weights", "weight_scale"],
+        output_names=["out"],
+        source=source,
+    )
+
+
+_bitlinear_kernel = make_bitlinear_kernel()
+
+
+class BitLinear(nn.Module):
+    """
+    BitLinear module with memory-efficient weight handling.
+    """
+
+    def __init__(
+        self,
+        in_features,
+        out_features,
+        bias=True,
+        invert_weight_scales=False,
+    ):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        # Calculate packed dimensions - the first dimension gets packed 4:1
+        # The weights are ternary so can be represented with 2 bits, and they
+        # are packed in uint8 tensors, hence the number of values per item is 4
+        packed_out_features = (out_features + 3) // 4
+        self.weight = mx.zeros((packed_out_features, in_features), dtype=mx.uint8)
+
+        self.invert_weight_scales = invert_weight_scales
+        self.weight_scale = mx.array([1.0])
+
+        if bias:
+            self.bias = mx.zeros((out_features,))
+        else:
+            self.bias = None
+
+    def execute_matmul_kernel(self, x, packed_weights):
+        original_shape = x.shape
+        if len(original_shape) > 2:
+            x = x.reshape(-1, original_shape[-1])
+        total_batch_elements, in_features = x.shape
+
+        out_features = self.out_features
+
+        dtype = self.weight_scale.dtype
+        assert x.dtype == dtype, "Wrong type for input."
+        out = _bitlinear_kernel(
+            inputs=[
+                x,
+                packed_weights,
+                self.weight_scale,
+            ],
+            template=[
+                ("T", dtype),
+                ("invert_weight_scales", self.invert_weight_scales),
+                ("in_features", in_features),
+                ("out_features", out_features),
+            ],
+            grid=(32, total_batch_elements * out_features // 4, 1),
+            threadgroup=(32, 1, 1),  # SIMD width is 32 threads
+            output_shapes=[(total_batch_elements, out_features)],
+            output_dtypes=[dtype],
+        )[0]
+
+        if len(original_shape) > 2:
+            out = out.reshape(*original_shape[:-1], out_features)
+        return out
+
+    def __call__(self, x):
+        y = self.execute_matmul_kernel(x, self.weight)
+
+        if self.bias is not None:
+            y = mx.add(y, self.bias)
+        return y
@@ -0,0 +1,215 @@
+# Copyright © 2023-2024 Apple Inc.
+
+from dataclasses import dataclass
+from functools import partial
+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 .bitlinear_layers import BitLinear
+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
+    num_key_value_heads: int
+    rms_norm_eps: float
+    vocab_size: int
+    head_dim: Optional[int] = None
+    max_position_embeddings: Optional[int] = None
+    attention_bias: bool = False
+    mlp_bias: bool = False
+    rope_theta: float = 10000
+    rope_traditional: bool = False
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = True
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+
+        self.head_dim = head_dim = args.head_dim or args.hidden_size // n_heads
+
+        self.scale = head_dim**-0.5
+        attention_bias = args.attention_bias
+
+        self.q_proj = BitLinear(dim, n_heads * head_dim, bias=attention_bias)
+        self.k_proj = BitLinear(dim, n_kv_heads * head_dim, bias=attention_bias)
+        self.v_proj = BitLinear(dim, n_kv_heads * head_dim, bias=attention_bias)
+        self.o_proj = BitLinear(n_heads * head_dim, dim, bias=attention_bias)
+
+        self.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            args.rope_traditional,
+            args.rope_scaling,
+            args.max_position_embeddings,
+        )
+        self.attn_sub_norm = 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:
+        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 = 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)
+        output = self.attn_sub_norm(output)
+        output = self.o_proj(output)
+
+        return output
+
+
+@partial(mx.compile, shapeless=True)
+def relu2(x):
+    return mx.square(nn.relu(x))
+
+
+class MLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        hidden_dim = args.intermediate_size
+        if hasattr(args, "mlp_bias"):
+            mlp_bias = args.mlp_bias
+        else:
+            mlp_bias = False
+
+        self.gate_proj = BitLinear(dim, hidden_dim, bias=mlp_bias)
+        self.down_proj = BitLinear(hidden_dim, dim, bias=mlp_bias)
+        self.up_proj = BitLinear(dim, hidden_dim, bias=mlp_bias)
+        self.ffn_sub_norm = nn.RMSNorm(args.intermediate_size, eps=args.rms_norm_eps)
+
+    def __call__(self, x) -> mx.array:
+        x = relu2(self.gate_proj(x)) * self.up_proj(x)
+        x = self.ffn_sub_norm(x)
+        x = self.down_proj(x)
+        return 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)
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        out = h + r
+        return out
+
+
+class LlamaModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.num_hidden_layers = args.num_hidden_layers
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            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, 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 = LlamaModel(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):
+        # Remove unused precomputed rotary freqs
+        weights = {
+            k: v for k, v in weights.items() if "self_attn.rotary_emb.inv_freq" not in k
+        }
+        if self.args.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
+        return weights
+
+    @property
+    def layers(self):
+        return self.model.layers
@@ -12,7 +12,7 @@ def make_prompt_cache(
    max_kv_size: Optional[int] = None,
 ) -> List[Any]:
    """
-    Construct the model's cache for use when cgeneration.
+    Construct the model's cache for use in generation.

    This function will defer the cache construction to the model if it has a
    ``make_cache`` method, otherwise it will make a default KV cache.
@@ -129,6 +129,40 @@ class _BaseCache:
        return False


+class ConcatenateKVCache(_BaseCache):
+    """ConcatenateKVCache the simplest KV cache implementation.
+
+    Can be used as a mock KV cache or when large blocks are being processed at
+    a time in which case KVCache isn't necessarily faster. Consider using the
+    KVCache with a larger step size before using this cache.
+    """
+
+    def __init__(self):
+        self.keys = None
+        self.values = None
+        self.offset = 0
+
+    def update_and_fetch(self, keys, values):
+        if self.keys is None:
+            self.keys = keys
+            self.values = values
+        else:
+            self.keys = mx.concatenate([self.keys, keys], axis=-2)
+            self.values = mx.concatenate([self.values, values], axis=-2)
+        self.offset = self.keys.shape[-2]
+
+        return self.keys, self.values
+
+    @property
+    def state(self):
+        return self.keys, self.values
+
+    @state.setter
+    def state(self, v):
+        self.keys, self.values = v
+        self.offset = self.keys.shape[-2]
+
+
 class QuantizedKVCache(_BaseCache):
    def __init__(self, group_size: int = 64, bits: int = 8):
        self.keys = None
@@ -1,7 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.

 from dataclasses import dataclass
-from typing import Any, Optional, Tuple
+from typing import Any, Optional

 import mlx.core as mx
 import mlx.nn as nn
@@ -1,7 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.

 from dataclasses import dataclass
-from typing import Any, Optional, Tuple
+from typing import Any, Optional

 import mlx.core as mx
 import mlx.nn as nn
@@ -105,10 +105,9 @@ class MLP(nn.Module):
        self.v1 = nn.Linear(d_model, ffn_dim, bias=False)
        self.w1 = nn.Linear(d_model, ffn_dim, bias=False)
        self.w2 = nn.Linear(ffn_dim, d_model, bias=False)
-        self.act_fn = nn.silu

    def __call__(self, x: mx.array) -> mx.array:
-        current_hidden_states = self.act_fn(self.w1(x)) * self.v1(x)
+        current_hidden_states = nn.silu(self.w1(x)) * self.v1(x)
        current_hidden_states = self.w2(current_hidden_states)
        return current_hidden_states

@@ -118,10 +118,9 @@ class DeepseekMLP(nn.Module):
        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
-        self.act_fn = nn.silu

    def __call__(self, x: mx.array) -> mx.array:
-        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))


 class MoEGate(nn.Module):
@@ -2,7 +2,7 @@

 import math
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, Optional

 import mlx.core as mx
 import mlx.nn as nn
@@ -3,7 +3,7 @@
 import math
 from dataclasses import dataclass
 from functools import partial
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, Optional

 import mlx.core as mx
 import mlx.nn as nn
@@ -130,6 +130,14 @@ def clipped_silu(x):
    return mx.clip(x * mx.sigmoid(x), a_min=-100, a_max=100)


+class ClippedSilu(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def __call__(self, x):
+        return clipped_silu(x)
+
+
 class DeepseekV3Attention(nn.Module):
    def __init__(self, config: ModelArgs):
        super().__init__()
@@ -344,7 +352,7 @@ class DeepseekV3MoE(nn.Module):
            config.hidden_size,
            config.moe_intermediate_size,
            config.n_routed_experts,
-            activation=clipped_silu,
+            activation=ClippedSilu(),
        )

        self.gate = MoEGate(config)
@@ -529,6 +537,7 @@ class Model(nn.Module):
    def layers(self):
        return self.model.layers[self.model.start_idx : self.model.end_idx]

+    @property
    def cast_predicate(self):
        def predicate(k):
            return "e_score_correction_bias" not in k
@@ -0,0 +1,320 @@
+# Copyright © 2023-2024 Apple Inc.
+
+from dataclasses import dataclass
+from functools import partial
+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
+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
+    rms_norm_eps: float
+    vocab_size: int
+    max_position_embeddings: Optional[int]
+    num_key_value_heads: Optional[int]
+    first_k_dense_replace: int
+    moe_intermediate_size: int
+    moe_layer_freq: int
+    n_routed_experts: int
+    n_shared_experts: int
+    norm_topk_prob: bool
+    num_experts_per_tok: int
+    rope_theta: float
+    routed_scaling_factor: float
+    head_dim: Optional[int] = None
+    scoring_func: str = ("noaux_tc",)
+    n_group: Optional[int] = 1
+    topk_group: Optional[int] = 1
+    attention_bias: bool = False
+    mlp_bias: bool = False
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = False
+
+
+class Dots1Attention(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 or args.hidden_size // n_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 = 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)
+
+
+@mx.compile
+def group_expert_select(
+    gates,
+    e_score_correction_bias,
+    top_k,
+    n_group,
+    topk_group,
+    routed_scaling_factor,
+    norm_topk_prob,
+):
+
+    k = top_k
+    scores = mx.sigmoid(gates.astype(mx.float32))
+    orig_scores = scores
+    scores = scores + e_score_correction_bias
+    k = n_group - topk_group
+    if k != 0:
+        scores = mx.unflatten(scores, axis=-1, shape=(n_group, -1))
+        group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True)
+        group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :]
+        scores = mx.put_along_axis(scores, group_idx, mx.array(0.0), axis=-2)
+        scores = mx.flatten(scores, -2, -1)
+
+    k = top_k
+    inds = mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k]
+    scores = mx.take_along_axis(orig_scores, inds, axis=-1)
+    if top_k > 1 and norm_topk_prob:
+        denominator = scores.sum(axis=-1, keepdims=True)
+        scores = scores / denominator
+    scores = scores * routed_scaling_factor
+
+    return inds, scores
+
+
+class Dots1TopkRouter(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.top_k = args.num_experts_per_tok
+        self.norm_topk_prob = args.norm_topk_prob
+        self.n_routed_experts = args.n_routed_experts
+        self.routed_scaling_factor = args.routed_scaling_factor
+        self.n_group = args.n_group
+        self.topk_group = args.topk_group
+        self.weight = mx.zeros((self.n_routed_experts, args.hidden_size))
+        self.e_score_correction_bias = mx.zeros((self.n_routed_experts,))
+
+    def __call__(self, x):
+        return group_expert_select(
+            x @ self.weight.T,
+            self.e_score_correction_bias,
+            self.top_k,
+            self.n_group,
+            self.topk_group,
+            self.routed_scaling_factor,
+            self.norm_topk_prob,
+        )
+
+
+class Dots1MLP(nn.Module):
+    def __init__(
+        self, args: ModelArgs, hidden_size: int = None, intermediate_size: int = None
+    ):
+        super().__init__()
+
+        self.hidden_size = args.hidden_size if hidden_size is None else hidden_size
+        self.intermediate_size = (
+            args.intermediate_size if intermediate_size is None else intermediate_size
+        )
+
+        self.gate_proj = nn.Linear(
+            self.hidden_size, self.intermediate_size, bias=args.mlp_bias
+        )
+        self.up_proj = nn.Linear(
+            self.hidden_size, self.intermediate_size, bias=args.mlp_bias
+        )
+        self.down_proj = nn.Linear(
+            self.intermediate_size, self.hidden_size, bias=args.mlp_bias
+        )
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class Dots1MoE(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_experts_per_tok = args.num_experts_per_tok
+        self.n_shared_experts = args.n_shared_experts
+
+        self.experts = SwitchGLU(
+            args.hidden_size,
+            args.moe_intermediate_size,
+            args.n_routed_experts,
+        )
+
+        self.gate = Dots1TopkRouter(args)
+
+        self.shared_experts = Dots1MLP(
+            args=args,
+            intermediate_size=args.moe_intermediate_size * args.n_shared_experts,
+        )
+
+    def __call__(self, x):
+        inds, scores = self.gate(x)
+        y = self.experts(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
+        if self.n_shared_experts is not None:
+            y = y + self.shared_experts(x)
+
+        return y
+
+
+class Dots1DecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.self_attn = Dots1Attention(args)
+
+        if layer_idx >= args.first_k_dense_replace:
+            self.mlp = Dots1MoE(args)
+        else:
+            self.mlp = Dots1MLP(args)
+
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        return h + r
+
+
+class Dots1Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            Dots1DecoderLayer(args, layer_idx)
+            for layer_idx 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,
+    ) -> mx.array:
+        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 = Dots1Model(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)
+
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}"
+            if l >= self.args.first_k_dense_replace:
+                for n, m in [
+                    ("w1", "gate_proj"),
+                    ("w2", "down_proj"),
+                    ("w3", "up_proj"),
+                ]:
+                    for k in ["weight", "scales", "biases"]:
+                        if f"{prefix}.mlp.experts.0.{m}.{k}" in weights:
+                            to_join = [
+                                weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}")
+                                for e in range(self.args.n_routed_experts)
+                            ]
+                            weights[f"{prefix}.mlp.experts.{m}.{k}"] = mx.stack(to_join)
+
+        return {k: v for k, v in weights.items() if "rotary_emb.inv_freq" not in k}
+
+    @property
+    def layers(self):
+        return self.model.layers
@@ -0,0 +1,167 @@
+# Copyright © 2023-2024 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
+from .rope_utils import initialize_rope
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    hidden_size: int
+    intermediate_size: int
+    model_type: str
+    max_position_embeddings: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    head_dim: Optional[int]
+    num_hidden_layers: int
+    rms_norm_eps: float
+    vocab_size: int
+    rope_theta: float
+    use_bias: bool
+    tie_word_embeddings: bool
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+
+        self.head_dim = head_dim = args.head_dim or dim // n_heads
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=args.use_bias)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.use_bias)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.use_bias)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=args.use_bias)
+
+        self.rope = initialize_rope(
+            head_dim,
+            base=args.rope_theta,
+            traditional=True,
+            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, use_bias=False):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=use_bias)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=use_bias)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=use_bias)
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.self_attn = Attention(args)
+        self.mlp = MLP(args.hidden_size, args.intermediate_size, args.use_bias)
+
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        return h + r
+
+
+class Ernie45Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [DecoderLayer(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 = Ernie45Model(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
+
+    @property
+    def layers(self):
+        return self.model.layers
@@ -1,7 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.

 from dataclasses import dataclass
-from typing import Any, Optional, Tuple
+from typing import Any, Optional

 import mlx.core as mx
 import mlx.nn as nn
@@ -1,7 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.

 from dataclasses import dataclass
-from typing import Any, Optional, Tuple
+from typing import Any, Optional

 import mlx.core as mx
 import mlx.nn as nn
@@ -41,8 +41,11 @@ class Model(nn.Module):
        inputs: mx.array,
        cache=None,
        mask: Optional[mx.array] = None,
+        input_embeddings: Optional[mx.array] = None,
    ):
-        return self.language_model(inputs, cache=cache, mask=mask)
+        return self.language_model(
+            inputs, cache=cache, mask=mask, input_embeddings=input_embeddings
+        )

    def sanitize(self, weights):
        weights = tree_unflatten(list(weights.items()))
@@ -175,9 +175,12 @@ class Gemma3Model(nn.Module):
        inputs: mx.array,
        mask: mx.array = None,
        cache=None,
+        input_embeddings: Optional[mx.array] = None,
    ):
-
-        h = self.embed_tokens(inputs)
+        if input_embeddings is not None:
+            h = input_embeddings
+        else:
+            h = self.embed_tokens(inputs)
        h *= mx.array(self.args.hidden_size**0.5, mx.bfloat16).astype(h.dtype)

        if cache is None:
@@ -218,8 +221,9 @@ class Model(nn.Module):
        inputs: mx.array,
        cache=None,
        mask: Optional[mx.array] = None,
+        input_embeddings: Optional[mx.array] = None,
    ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, mask, cache, input_embeddings)
        out = self.lm_head(out)
        return out

@@ -0,0 +1,621 @@
+# Copyright © 2025 Apple Inc.
+
+import math
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, Dict, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten, tree_unflatten
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import KVCache, RotatingKVCache
+
+
+@dataclass
+class TextConfig(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    head_dim: int
+    rms_norm_eps: float
+    vocab_size: int
+    num_key_value_heads: int
+    num_kv_shared_layers: int
+    query_pre_attn_scalar: float
+    vocab_size_per_layer_input: int
+    sliding_window: int
+    max_position_embeddings: int
+    rope_local_base_freq: float
+    rope_theta: float
+    final_logit_softcapping: float
+    layer_types: List[str]
+    activation_sparsity_pattern: List[float]
+    hidden_size_per_layer_input: int
+    altup_num_inputs: int
+    altup_coef_clip: float
+    altup_correct_scale: bool
+    altup_active_idx: int
+    laurel_rank: int
+    rope_scaling: Optional[Dict] = None
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    text_config: dict
+
+
+class RMSNoScale(nn.Module):
+    def __init__(self, eps: float = 1e-5):
+        super().__init__()
+        self.eps = eps
+
+    def __call__(self, x):
+        return mx.fast.rms_norm(x, None, self.eps)
+
+
+class Gemma3nLaurelBlock(nn.Module):
+    """Learned Augmented Residual Layer"""
+
+    def __init__(self, config: TextConfig):
+        super().__init__()
+        self.config = config
+
+        self.linear_left = nn.Linear(
+            self.config.hidden_size, self.config.laurel_rank, bias=False
+        )
+        self.linear_right = nn.Linear(
+            self.config.laurel_rank, self.config.hidden_size, bias=False
+        )
+        self.post_laurel_norm = nn.RMSNorm(
+            dims=self.config.hidden_size,
+            eps=self.config.rms_norm_eps,
+        )
+
+    def __call__(self, x: mx.array) -> mx.array:
+        laurel_x = self.linear_left(x)
+        laurel_x = self.linear_right(laurel_x)
+        normed_laurel_x = self.post_laurel_norm(laurel_x)
+        return x + normed_laurel_x
+
+
+class Gemma3nAttention(nn.Module):
+    def __init__(self, config: TextConfig, layer_idx: int, is_kv_shared_layer: bool):
+        super().__init__()
+        self.is_sliding = config.layer_types[layer_idx] == "sliding_attention"
+
+        dim = config.hidden_size
+        self.n_heads = n_heads = config.num_attention_heads
+        self.n_kv_heads = n_kv_heads = config.num_key_value_heads
+        self.repeats = n_heads // n_kv_heads
+        self.head_dim = head_dim = config.head_dim
+        self.layer_idx = layer_idx
+
+        self.scale = 1.0
+
+        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(dims=config.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = nn.RMSNorm(dims=config.head_dim, eps=config.rms_norm_eps)
+        self.v_norm = RMSNoScale(eps=config.rms_norm_eps)
+
+        self.is_kv_shared_layer = is_kv_shared_layer
+
+        self.rope = nn.RoPE(
+            head_dim,
+            traditional=False,
+            base=(
+                config.rope_local_base_freq if self.is_sliding else config.rope_theta
+            ),
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, _ = x.shape
+
+        queries = self.q_proj(x)
+        queries = queries.reshape(B, L, -1, self.head_dim)
+        queries = self.q_norm(queries)
+
+        offset = 0
+        if self.is_kv_shared_layer and cache is not None:
+            # For shared layers, retrieve KV from the designated cache layer
+            keys, values = cache.state
+            offset = cache.offset
+
+        else:
+            if cache is not None:
+                offset = cache.offset
+            keys = self.k_proj(x).reshape(B, L, -1, self.head_dim)
+            keys = self.k_norm(keys)
+            keys = keys.transpose(0, 2, 1, 3)
+            keys = self.rope(keys, offset=offset)
+
+            values = self.v_proj(x).reshape(B, L, -1, self.head_dim)
+            values = self.v_norm(values)
+            values = values.transpose(0, 2, 1, 3)
+
+            if cache is not None:
+                keys, values = cache.update_and_fetch(keys, values)
+
+        queries = queries.transpose(0, 2, 1, 3)
+        queries = self.rope(queries, offset=offset)
+
+        if isinstance(mask, mx.array) and mask.shape[-1] != keys.shape[-2]:
+            mask = mask[:, : keys.shape[-2]]
+
+        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)
+
+
+@partial(mx.compile, shapeless=True)
+def gelu_topk(inputs, std_multiplier):
+    inputs_mean = mx.mean(inputs, axis=-1, keepdims=True)
+    inputs_std = mx.std(inputs, axis=-1, keepdims=True)
+    cutoff_x = inputs_mean + inputs_std * std_multiplier.astype(inputs_std.dtype)
+    return nn.gelu_approx(mx.maximum(0, inputs - cutoff_x))
+
+
+class MLP(nn.Module):
+    def __init__(self, config: TextConfig, layer_idx: int = 0):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        if config.activation_sparsity_pattern is not None:
+            self.activation_sparsity = config.activation_sparsity_pattern[layer_idx]
+        else:
+            self.activation_sparsity = 0.0
+        if self.activation_sparsity > 0:
+            self._std_multiplier = math.sqrt(2.0) * mx.erfinv(
+                2 * self.activation_sparsity - 1
+            )
+
+    def __call__(self, x: mx.array):
+        gate_proj = self.gate_proj(x)
+        if self.activation_sparsity > 0.0:
+            activations = gelu_topk(gate_proj, self._std_multiplier)
+        else:
+            activations = nn.gelu_approx(gate_proj)
+        up_proj = self.up_proj(x)
+        down_proj = self.down_proj(activations * up_proj)
+        return down_proj
+
+
+class Gemma3nAltUp(nn.Module):
+    """Alternating Updates (AltUp)"""
+
+    def __init__(self, config: TextConfig):
+        super().__init__()
+        self.config = config
+
+        self.correct_output_scale = mx.zeros((self.config.hidden_size,))
+        self.correction_coefs = nn.Linear(
+            self.config.altup_num_inputs, self.config.altup_num_inputs, bias=False
+        )
+        self.prediction_coefs = nn.Linear(
+            self.config.altup_num_inputs, self.config.altup_num_inputs**2, bias=False
+        )
+        self.modality_router = nn.Linear(
+            self.config.hidden_size, self.config.altup_num_inputs, bias=False
+        )
+        self.router_norm = nn.RMSNorm(
+            dims=self.config.hidden_size,
+            eps=self.config.rms_norm_eps,
+        )
+
+    def compute_router_modalities(self, x: mx.array) -> mx.array:
+        router_inputs = self.router_norm(x) * (self.config.hidden_size**-1.0)
+        routed = self.modality_router(router_inputs).astype(mx.float32)
+        return mx.tanh(routed)
+
+    def predict(self, x: mx.array) -> mx.array:
+        modalities = self.compute_router_modalities(x[self.config.altup_active_idx])
+
+        self.prediction_coefs.weight = self.prediction_coefs.weight.astype(mx.float32)
+
+        if self.config.altup_coef_clip is not None:
+            self.prediction_coefs.weight = mx.clip(
+                self.prediction_coefs.weight,
+                -self.config.altup_coef_clip,
+                self.config.altup_coef_clip,
+            )
+
+        all_coefs = (
+            self.prediction_coefs(modalities)
+            .reshape(
+                *modalities.shape[:-1],
+                self.config.altup_num_inputs,
+                self.config.altup_num_inputs,
+            )
+            .transpose(0, 1, 3, 2)
+        )
+
+        x_up = x.astype(mx.float32)
+        x_permuted = x_up.transpose(1, 2, 3, 0)
+        predictions = mx.matmul(x_permuted, all_coefs)
+        predictions = predictions.transpose(3, 0, 1, 2)
+        predictions += x_up
+        return predictions.astype(x.dtype)
+
+    def correct(self, predictions: mx.array, activated: mx.array):
+        modalities = self.compute_router_modalities(activated)
+
+        self.correction_coefs.weight = self.correction_coefs.weight.astype(mx.float32)
+
+        if self.config.altup_coef_clip is not None:
+            self.correction_coefs.weight = mx.clip(
+                self.correction_coefs.weight,
+                -self.config.altup_coef_clip,
+                self.config.altup_coef_clip,
+            )
+
+        all_coefs = self.correction_coefs(modalities) + 1.0
+
+        active_x = predictions[self.config.altup_active_idx]
+        innovation = activated - active_x
+
+        all_coefs = all_coefs.transpose(2, 1, 0)
+        corrected = innovation[None] * all_coefs[:, None]
+        corrected += predictions
+
+        return corrected.astype(activated.dtype)
+
+
+class Gemma3nDecoderLayer(nn.Module):
+    def __init__(self, config: TextConfig, layer_idx: int, is_kv_shared_layer: bool):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.layer_idx = layer_idx
+        self.self_attn = Gemma3nAttention(config, layer_idx, is_kv_shared_layer)
+        self.mlp = MLP(config, layer_idx=layer_idx)
+        self.input_layernorm = nn.RMSNorm(
+            self.hidden_size,
+            eps=config.rms_norm_eps,
+        )
+
+        self.post_attention_layernorm = nn.RMSNorm(
+            self.hidden_size,
+            eps=config.rms_norm_eps,
+        )
+        self.pre_feedforward_layernorm = nn.RMSNorm(
+            self.hidden_size,
+            eps=config.rms_norm_eps,
+        )
+        self.post_feedforward_layernorm = nn.RMSNorm(
+            self.hidden_size,
+            eps=config.rms_norm_eps,
+        )
+        self.is_sliding = self.self_attn.is_sliding
+        self.sliding_window = config.sliding_window
+
+        self.hidden_size_per_layer_input = config.hidden_size_per_layer_input
+
+        self.altup = Gemma3nAltUp(config)
+        self.laurel = Gemma3nLaurelBlock(config)
+        self.per_layer_input_gate = nn.Linear(
+            self.hidden_size, self.hidden_size_per_layer_input, bias=False
+        )
+        self.per_layer_projection = nn.Linear(
+            self.hidden_size_per_layer_input, self.hidden_size, bias=False
+        )
+        self.post_per_layer_input_norm = nn.RMSNorm(
+            self.hidden_size,
+            eps=config.rms_norm_eps,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+        per_layer_input: Optional[mx.array] = None,
+    ):
+        predictions = self.altup.predict(x)
+        active_prediction = predictions[self.config.altup_active_idx]
+
+        active_prediction_normed = self.input_layernorm(active_prediction)
+        laurel_output = self.laurel(active_prediction_normed)
+
+        attn = self.self_attn(
+            active_prediction_normed,
+            mask,
+            cache,
+        )
+
+        attn = self.post_attention_layernorm(attn)
+
+        attn_gated = active_prediction + attn
+        attn_laurel = (attn_gated + laurel_output) * (2.0**-0.5)
+
+        attn_norm = self.pre_feedforward_layernorm(attn_laurel)
+        attn_ffw = self.mlp(attn_norm)
+        attn_ffw_norm = self.post_feedforward_layernorm(attn_ffw)
+        attn_ffw_laurel_gated = attn_laurel + attn_ffw_norm
+
+        corrected_predictions = self.altup.correct(predictions, attn_ffw_laurel_gated)
+
+        first_prediction = corrected_predictions[self.config.altup_active_idx]
+        if self.config.altup_correct_scale:
+            first_prediction = first_prediction * self.altup.correct_output_scale
+
+        first_prediction = self.per_layer_input_gate(first_prediction)
+        first_prediction = nn.gelu_approx(first_prediction)
+
+        first_prediction = mx.multiply(first_prediction, per_layer_input)
+
+        first_prediction = self.per_layer_projection(first_prediction)
+        first_prediction = self.post_per_layer_input_norm(first_prediction)
+
+        corrected_predictions[1:] = corrected_predictions[1:] + first_prediction
+
+        return corrected_predictions
+
+
+@partial(mx.compile, shapeless=True)
+def logit_softcap(softcap, x):
+    out = mx.tanh(x / softcap)
+    out = out * softcap
+    return out
+
+
+class LanguageModel(nn.Module):
+    def __init__(self, config: TextConfig):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.hidden_size_per_layer_input = config.hidden_size_per_layer_input
+        self.vocab_size = config.vocab_size
+        self.vocab_size_per_layer_input = config.vocab_size_per_layer_input
+        self.num_hidden_layers = config.num_hidden_layers
+        self.final_logit_softcapping = config.final_logit_softcapping
+        self.first_kv_shared_layer_idx = (
+            config.num_hidden_layers - config.num_kv_shared_layers
+        )
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = [
+            Gemma3nDecoderLayer(
+                config=config,
+                layer_idx=layer_idx,
+                is_kv_shared_layer=layer_idx >= self.first_kv_shared_layer_idx,
+            )
+            for layer_idx in range(config.num_hidden_layers)
+        ]
+
+        self.embed_tokens_per_layer = nn.Embedding(
+            config.vocab_size_per_layer_input,
+            config.num_hidden_layers * config.hidden_size_per_layer_input,
+        )
+
+        self.per_layer_model_projection = nn.Linear(
+            config.hidden_size,
+            config.num_hidden_layers * config.hidden_size_per_layer_input,
+            bias=False,
+        )
+
+        self.per_layer_projection_norm = nn.RMSNorm(
+            dims=config.hidden_size_per_layer_input,
+            eps=config.rms_norm_eps,
+        )
+
+        self.altup_projections = [
+            nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+            for _ in range(1, self.config.altup_num_inputs)
+        ]
+
+        self.altup_unembed_projections = [
+            nn.Linear(config.hidden_size, config.hidden_size, bias=False)
+            for _ in range(1, self.config.altup_num_inputs)
+        ]
+
+        self.norm = nn.RMSNorm(
+            config.hidden_size,
+            eps=config.rms_norm_eps,
+        )
+
+        self.first_sliding_idx = self.config.layer_types.index("sliding_attention")
+        self.first_full_idx = self.config.layer_types.index("full_attention")
+
+        concrete_layers = self.config.layer_types[: self.first_kv_shared_layer_idx]
+        shared_full_idx = (
+            len(concrete_layers) - 1 - concrete_layers[::-1].index("full_attention")
+        )
+        shared_sliding_idx = (
+            len(concrete_layers) - 1 - concrete_layers[::-1].index("sliding_attention")
+        )
+
+        self.layer_idx_to_cache_idx = []
+        for i, layer_type in enumerate(self.config.layer_types):
+            if i < self.first_kv_shared_layer_idx:
+                self.layer_idx_to_cache_idx.append(i)
+            else:
+                if layer_type == "full_attention":
+                    self.layer_idx_to_cache_idx.append(shared_full_idx)
+                elif layer_type == "sliding_attention":
+                    self.layer_idx_to_cache_idx.append(shared_sliding_idx)
+                else:
+                    raise NotImplementedError(f"Unknown layer type: {layer_type}")
+
+    def __call__(
+        self,
+        inputs: mx.array = None,
+        mask: mx.array = None,
+        cache=None,
+        input_embeddings: mx.array = None,
+    ):
+        if input_embeddings is None:
+            h = self.embed_tokens(inputs) * (self.hidden_size**0.5)
+        else:
+            h = input_embeddings
+
+        per_layer_inputs = self.get_per_layer_inputs(inputs)
+        per_layer_inputs = self.project_per_layer_inputs(h, per_layer_inputs)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        if mask is None:
+            full_mask = create_attention_mask(
+                h,
+                cache[self.first_full_idx :],
+            )
+            sliding_window_mask = create_attention_mask(
+                h,
+                cache[self.first_sliding_idx :],
+            )
+        h0 = h
+
+        # Expand hidden_states to support per-layer inputs
+        target_magnitude = mx.mean(h0**2, axis=-1, keepdims=True) ** 0.5
+
+        h_list = [h0]
+        h_list.extend([proj(h0) for proj in self.altup_projections])
+        h = mx.stack(h_list, axis=0)
+        mags = mx.mean(h[1:] ** 2, axis=-1, keepdims=True) ** 0.5
+        h[1:] = h[1:] * (target_magnitude / mx.maximum(mags, mx.finfo(h0.dtype).min))
+
+        for i, layer in enumerate(self.layers):
+            per_layer_input = per_layer_inputs[:, :, i, :]
+
+            is_global = self.config.layer_types[i] == "full_attention"
+
+            local_mask = mask
+            if mask is None and is_global:
+                local_mask = full_mask
+            elif mask is None:
+                local_mask = sliding_window_mask
+
+            h = layer(
+                h,
+                local_mask,
+                cache[self.layer_idx_to_cache_idx[i]],
+                per_layer_input,
+            )
+
+        # Per-layer inputs to single output
+        target_magnitude = mx.mean(h[0] ** 2, axis=-1, keepdims=True) ** 0.5
+        for i, proj in enumerate(self.altup_unembed_projections):
+            h[i + 1] = proj(h[i + 1])
+        mags = mx.mean(h[1:] ** 2, axis=-1, keepdims=True) ** 0.5
+        h[1:] = h[1:] * (target_magnitude / mx.maximum(mags, mx.finfo(h0.dtype).min))
+
+        h = mx.mean(h, axis=0)
+
+        out = self.norm(h)
+        out = self.embed_tokens.as_linear(out)
+        if self.final_logit_softcapping is not None:
+            out = logit_softcap(self.final_logit_softcapping, out)
+        return out
+
+    def get_per_layer_inputs(self, input_ids: mx.array) -> mx.array:
+        per_layer_inputs_mask = input_ids < self.vocab_size_per_layer_input
+        tokens = mx.where(per_layer_inputs_mask, input_ids, mx.zeros_like(input_ids))
+        result = self.embed_tokens_per_layer(tokens) * (
+            self.hidden_size_per_layer_input**0.5
+        )
+        return result.reshape(
+            *input_ids.shape,
+            self.num_hidden_layers,
+            self.hidden_size_per_layer_input,
+        )
+
+    def project_per_layer_inputs(
+        self,
+        inputs_embeds: mx.array,
+        per_layer_inputs: mx.array,
+    ) -> mx.array:
+        per_layer_projection = self.per_layer_model_projection(inputs_embeds) * (
+            self.hidden_size**-0.5
+        )
+        per_layer_projection = per_layer_projection.reshape(
+            *inputs_embeds.shape[:-1],
+            self.config.num_hidden_layers,
+            self.config.hidden_size_per_layer_input,
+        )
+        per_layer_projection = self.per_layer_projection_norm(per_layer_projection)
+        return (per_layer_projection + per_layer_inputs) * (2.0**-0.5)
+
+    def make_cache(self):
+        caches = []
+        for layer_type in self.config.layer_types[: self.first_kv_shared_layer_idx]:
+            if layer_type == "full_attention":
+                caches.append(KVCache())
+            elif layer_type == "sliding_attention":
+                caches.append(
+                    RotatingKVCache(max_size=self.config.sliding_window, keep=0)
+                )
+            else:
+                raise NotImplementedError(f"Unknown layer type: {layer_type}")
+        return caches
+
+
+class Gemma3n(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.language_model = LanguageModel(TextConfig.from_dict(args.text_config))
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+        mask: Optional[mx.array] = None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        return self.language_model(
+            inputs, cache=cache, mask=mask, input_embeddings=input_embeddings
+        )
+
+    def make_cache(self):
+        return self.language_model.make_cache()
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model = Gemma3n(args)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+        mask: Optional[mx.array] = None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        return self.model(
+            inputs, cache=cache, mask=mask, input_embeddings=input_embeddings
+        )
+
+    def sanitize(self, weights):
+        weights = tree_unflatten(list(weights.items()))
+        for k in ["vision_tower", "audio_tower", "embed_audio", "embed_vision"]:
+            weights["model"].pop(k, None)
+        return dict(tree_flatten(weights))
+
+    @property
+    def layers(self):
+        return self.model.language_model.layers
+
+    def make_cache(self):
+        return self.model.make_cache()
@@ -5,7 +5,6 @@ from typing import Any, Optional

 import mlx.core as mx
 import mlx.nn as nn
-import numpy as np

 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention

@@ -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
@@ -1,11 +1,10 @@
 # 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
-import numpy as np

 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention

@@ -1,7 +1,7 @@
 # Copyright © 2025 Apple Inc.

 from dataclasses import dataclass
-from typing import Any, Optional, Tuple
+from typing import Any, Optional

 import mlx.core as mx
 import mlx.nn as nn
@@ -1,6 +1,5 @@
 # Copyright © 2023-2024 Apple Inc.

-import math
 from dataclasses import dataclass
 from typing import Any, Dict, Optional, Tuple, Union

@@ -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, Dict, Optional, Union

 import mlx.core as mx
 import mlx.nn as nn
@@ -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, Dict, Optional, Union

 import mlx.core as mx
 import mlx.nn as nn
@@ -1,16 +1,13 @@
 # Copyright © 2024 Apple Inc.

-import math
 from dataclasses import dataclass
-from functools import partial
-from typing import Any, Dict, Optional, Tuple, Union
+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 .base import BaseModelArgs
 from .deepseek_v3 import DeepseekV3Model
-from .switch_layers import SwitchGLU


@dataclass
@@ -113,6 +110,7 @@ class Model(nn.Module):
    def layers(self):
        return self.language_model.model.layers

+    @property
    def cast_predicate(self):
        def predicate(k):
            return "e_score_correction_bias" not in k
@@ -157,8 +157,12 @@ class LlamaModel(nn.Module):
        inputs: mx.array,
        mask: mx.array = None,
        cache=None,
+        input_embeddings: Optional[mx.array] = None,
    ):
-        h = self.embed_tokens(inputs)
+        if input_embeddings is not None:
+            h = input_embeddings
+        else:
+            h = self.embed_tokens(inputs)

        if mask is None:
            mask = create_attention_mask(h, cache)
@@ -186,8 +190,9 @@ class Model(nn.Module):
        inputs: mx.array,
        mask: mx.array = None,
        cache=None,
+        input_embeddings: Optional[mx.array] = None,
    ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, mask, cache, input_embeddings)
        if self.args.tie_word_embeddings:
            out = self.model.embed_tokens.as_linear(out)
        else:
@@ -1,7 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.

 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Union
+from typing import Any, Optional, Union

 import mlx.core as mx
 import mlx.nn as nn
@@ -1,3 +1,5 @@
+# Copyright © 2023-2025 Apple Inc.
+
 from dataclasses import dataclass
 from typing import Any, Dict, Optional, Union

@@ -7,6 +7,7 @@ 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
@@ -22,6 +23,7 @@ class ModelArgs(BaseModelArgs):
    num_key_value_heads: int
    scale_depth: float
    scale_emb: float
+    max_position_embeddings: Optional[int] = None
    rope_theta: float = 1000000.0
    rope_traditional: bool = False
    rope_scaling: Optional[Dict[str, Union[str, float]]] = None
@@ -67,17 +69,12 @@ class Attention(nn.Module):
            self.num_heads * self.head_dim, self.hidden_size, bias=False
        )

-        rope_scale = (
-            1 / args.rope_scaling["factor"]
-            if args.rope_scaling is not None and args.rope_scaling["type"] == "linear"
-            else 1
-        )
-
-        self.rope = nn.RoPE(
-            dims=self.head_dim,
-            traditional=args.rope_traditional,
-            base=self.rope_theta,
-            scale=rope_scale,
+        self.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            args.rope_traditional,
+            args.rope_scaling,
+            args.max_position_embeddings,
        )

    def __call__(
@@ -7,7 +7,7 @@ import mlx.core as mx
 import mlx.nn as nn

 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
-from .su_rope import SuScaledRotaryEmbedding
+from .rope_utils import SuScaledRoPE


@dataclass
@@ -82,7 +82,7 @@ class Attention(nn.Module):
            bias=self.attention_bias,
        )

-        self.rope = SuScaledRotaryEmbedding(
+        self.rope = SuScaledRoPE(
            dims=args.qk_rope_head_dim,
            base=args.rope_theta,
            max_position_embeddings=args.max_position_embeddings,
@@ -0,0 +1,49 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten, tree_unflatten
+
+from . import llama
+from .base import BaseModelArgs
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    text_config: dict
+
+    def __post_init__(self):
+        self.text_config["tie_word_embeddings"] = False
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.language_model = llama.Model(llama.ModelArgs.from_dict(args.text_config))
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+        mask: Optional[mx.array] = None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        return self.language_model(
+            inputs, cache=cache, mask=mask, input_embeddings=input_embeddings
+        )
+
+    def sanitize(self, weights):
+        weights = tree_unflatten(list(weights.items()))
+        weights.pop("vision_tower", None)
+        weights.pop("multi_modal_projector", None)
+        return dict(tree_flatten(weights))
+
+    @property
+    def layers(self):
+        return self.language_model.model.layers
@@ -1,8 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.

-import math
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Dict, Optional, Union

 import mlx.core as mx
 import mlx.nn as nn
@@ -1,7 +1,6 @@
 # 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
@@ -2,7 +2,7 @@

 import sys
 from dataclasses import dataclass
-from typing import Any, Optional, Tuple
+from typing import Any, Optional

 import mlx.core as mx
 import mlx.nn as nn
@@ -1,7 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.

 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Union

 import mlx.core as mx
 import mlx.nn as nn
@@ -2,7 +2,6 @@

 import math
 from dataclasses import dataclass
-from typing import Tuple

 import mlx.core as mx
 import mlx.nn as nn
@@ -112,10 +111,9 @@ class PhiMLP(nn.Module):
        super().__init__()
        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
-        self.act = nn.GELU(approx="precise")

    def __call__(self, x) -> mx.array:
-        return self.fc2(self.act(self.fc1(x)))
+        return self.fc2(nn.gelu_approx(self.fc1(x)))


 class PhiDecoderLayer(nn.Module):
@@ -7,7 +7,7 @@ import mlx.core as mx
 import mlx.nn as nn

 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
-from .su_rope import SuScaledRotaryEmbedding
+from .rope_utils import SuScaledRoPE


@dataclass
@@ -63,7 +63,7 @@ class Attention(nn.Module):

        rope_dim = int(head_dim * args.partial_rotary_factor)
        if args.rope_scaling and args.rope_scaling["type"] in ["longrope", "su"]:
-            self.rope = SuScaledRotaryEmbedding(
+            self.rope = SuScaledRoPE(
                rope_dim,
                base=args.rope_theta,
                max_position_embeddings=args.max_position_embeddings,
@@ -7,7 +7,7 @@ import mlx.core as mx
 import mlx.nn as nn

 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
-from .su_rope import SuScaledRotaryEmbedding
+from .rope_utils import SuScaledRoPE
 from .switch_layers import SwitchGLU


@@ -45,7 +45,7 @@ class Attention(nn.Module):
        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=True)
        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=True)

-        self.rope = SuScaledRotaryEmbedding(
+        self.rope = SuScaledRoPE(
            head_dim,
            base=args.rope_theta,
            max_position_embeddings=args.max_position_embeddings,
@@ -0,0 +1,52 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten, tree_unflatten
+
+from . import llama
+from .base import BaseModelArgs
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    text_config: dict
+
+    def __post_init__(self):
+        self.text_config["tie_word_embeddings"] = False
+        self.text_config["num_attention_heads"] = self.text_config.get(
+            "num_attention_heads", 32
+        )
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.language_model = llama.Model(llama.ModelArgs.from_dict(args.text_config))
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+        mask: Optional[mx.array] = None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        return self.language_model(
+            inputs, cache=cache, mask=mask, input_embeddings=input_embeddings
+        )
+
+    def sanitize(self, weights):
+        weights = tree_unflatten(list(weights.items()))
+        weights.pop("vision_tower", None)
+        weights.pop("multi_modal_projector", None)
+        return dict(tree_flatten(weights))
+
+    @property
+    def layers(self):
+        return self.language_model.model.layers
@@ -137,8 +137,12 @@ class Qwen2Model(nn.Module):
        inputs: mx.array,
        mask: mx.array = None,
        cache=None,
+        input_embeddings: Optional[mx.array] = None,
    ):
-        h = self.embed_tokens(inputs)
+        if input_embeddings is not None:
+            h = input_embeddings
+        else:
+            h = self.embed_tokens(inputs)

        if mask is None:
            mask = create_attention_mask(h, cache)
@@ -166,8 +170,9 @@ class Model(nn.Module):
        inputs: mx.array,
        mask: mx.array = None,
        cache=None,
+        input_embeddings: Optional[mx.array] = None,
    ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, mask, cache, input_embeddings)
        if self.args.tie_word_embeddings:
            out = self.model.embed_tokens.as_linear(out)
        else:
@@ -1,6 +1,5 @@
 # Copyright © 2023-2024 Apple Inc.

-import math
 from dataclasses import dataclass
 from typing import Any, Dict, Optional, Union

@@ -1,6 +1,5 @@
 # Copyright © 2025 Apple Inc.

-import math
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional, Union

@@ -8,7 +7,6 @@ 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


@@ -19,7 +17,6 @@ class ModelArgs(BaseModelArgs):
    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
@@ -2,7 +2,7 @@

 import math
 from dataclasses import dataclass
-from typing import List, Literal, Optional
+from typing import List, Optional

 import mlx.core as mx
 import mlx.nn as nn
@@ -1,12 +1,71 @@
 # Copyright © 2023-2024 Apple Inc.

 import math
-from typing import Optional
+from typing import List, Optional, Union

 import mlx.core as mx
 import mlx.nn as nn


+class SuScaledRoPE(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        base: float = 10000.0,
+        max_position_embeddings: int = 131072,
+        original_max_position_embeddings: int = 4096,
+        short_factor: Union[List[float], float] = 1.0,
+        long_factor: Union[List[float], float] = 1.0,
+        short_mscale: float = None,
+        long_mscale: float = None,
+    ):
+        """
+        Su Scaled Rotary Embedding layer.
+
+        Args:
+            dims (int): The feature dimensions to be rotated.
+            base (int, optional): Base for the exponential scaling.
+            max_position_embeddings (int, optional): The maximum sequence
+              length that this model was trained with. This is used to determine
+              the size of the original RoPE embeddings when using long scaling.
+              Default: ``131072``.
+            original_max_position_embeddings (int, optional): The maximum
+              sequence length that this model was trained with. This is used to
+              determine the size of the original RoPE embeddings when using long
+              scaling. Default: ``4096``.
+            short_factor (float or list[float], optional): List of scaling
+              factors for sequences of length lesser than
+              ``original_max_position_embeddings``. Default: ``1.0``.
+            long_factor (float or list[float], optional): List of scaling
+              factors for sequences of length greater than
+              ``original_max_position_embeddings``.  Default: ``1.0``.
+            short_mscale (float, optional): Scale the input prior to embedding.
+            long_mscale (float, optional): Scale the input prior to embedding.
+        """
+        super().__init__()
+        freqs = base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
+        self._freqs = mx.array(long_factor, dtype=mx.float32) * freqs
+        self.original_max_position_embeddings = original_max_position_embeddings
+        self.scale = long_mscale or math.sqrt(
+            1
+            + math.log(max_position_embeddings / original_max_position_embeddings)
+            / math.log(original_max_position_embeddings)
+        )
+        self.dim = dims
+
+    def __call__(self, x, offset: int = 0):
+        x[..., : self.dim] = self.scale * x[..., : self.dim]
+        return mx.fast.rope(
+            x,
+            self.dim,
+            traditional=False,
+            base=None,
+            scale=1.0,
+            offset=offset,
+            freqs=self._freqs,
+        )
+
+
 class Llama3RoPE(nn.Module):

    def __init__(
@@ -180,5 +239,17 @@ def initialize_rope(
            base=base,
            **rope_kwargs,
        )
+    elif rope_type == "longrope":
+        return SuScaledRoPE(
+            dims=dims,
+            base=base,
+            max_position_embeddings=max_position_embeddings,
+            original_max_position_embeddings=scaling_config[
+                "original_max_position_embeddings"
+            ],
+            short_factor=scaling_config["short_factor"],
+            long_factor=scaling_config["long_factor"],
+        )
+
    else:
        raise ValueError(f"Unsupported RoPE type {rope_type}")
@@ -1,66 +0,0 @@
-# Copyright © 2023-2024 Apple Inc.
-
-import math
-from typing import List, Union
-
-import mlx.core as mx
-import mlx.nn as nn
-
-
-class SuScaledRotaryEmbedding(nn.Module):
-    def __init__(
-        self,
-        dims: int,
-        base: float = 10000.0,
-        max_position_embeddings: int = 131072,
-        original_max_position_embeddings: int = 4096,
-        short_factor: Union[List[float], float] = 1.0,
-        long_factor: Union[List[float], float] = 1.0,
-        short_mscale: float = None,
-        long_mscale: float = None,
-    ):
-        """
-        Su Scaled Rotary Embedding layer.
-
-        Args:
-            dims (int): The feature dimensions to be rotated.
-            base (int, optional): Base for the exponential scaling.
-            max_position_embeddings (int, optional): The maximum sequence
-              length that this model was trained with. This is used to determine
-              the size of the original RoPE embeddings when using long scaling.
-              Default: ``131072``.
-            original_max_position_embeddings (int, optional): The maximum
-              sequence length that this model was trained with. This is used to
-              determine the size of the original RoPE embeddings when using long
-              scaling. Default: ``4096``.
-            short_factor (float or list[float], optional): List of scaling
-              factors for sequences of length lesser than
-              ``original_max_position_embeddings``. Default: ``1.0``.
-            long_factor (float or list[float], optional): List of scaling
-              factors for sequences of length greater than
-              ``original_max_position_embeddings``.  Default: ``1.0``.
-            short_mscale (float, optional): Scale the input prior to embedding.
-            long_mscale (float, optional): Scale the input prior to embedding.
-        """
-        super().__init__()
-        freqs = base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
-        self._freqs = mx.array(long_factor, dtype=mx.float32) * freqs
-        self.original_max_position_embeddings = original_max_position_embeddings
-        self.scale = long_mscale or math.sqrt(
-            1
-            + math.log(max_position_embeddings / original_max_position_embeddings)
-            / math.log(original_max_position_embeddings)
-        )
-        self.dim = dims
-
-    def __call__(self, x, offset: int = 0):
-        x[..., : self.dim] = self.scale * x[..., : self.dim]
-        return mx.fast.rope(
-            x,
-            self.dim,
-            traditional=False,
-            base=None,
-            scale=1.0,
-            offset=offset,
-            freqs=self._freqs,
-        )
@@ -53,12 +53,6 @@ class QuantizedSwitchLinear(nn.Module):
        # Freeze this model's parameters
        self.freeze()

-    def unfreeze(self, *args, **kwargs):
-        """Wrap unfreeze so that we unfreeze any layers we might contain but
-        our parameters will remain frozen."""
-        super().unfreeze(*args, **kwargs)
-        self.freeze(recurse=False)
-
    @property
    def input_dims(self):
        return self.scales.shape[2] * self.group_size
@@ -14,6 +14,7 @@ from tqdm import tqdm

 from mlx_lm.models.base import create_attention_mask
 from mlx_lm.models.switch_layers import SwitchLinear
+from mlx_lm.quant.utils import load_data
 from mlx_lm.utils import (
    fetch_from_hub,
    get_model_path,
@@ -510,23 +511,6 @@ def awq_quantize(
        )


-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],
@@ -571,14 +555,14 @@ def main():

    mx.random.seed(args.seed)

-    model_path = get_model_path(args.model, revision=None)
+    model_path, hf_repo = 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 = load_data(tokenizer, args.num_samples, args.sequence_length)

    calibration_data = dist_split(calibration_data, group)

@@ -594,12 +578,11 @@ def main():
    )

    config = update_config(model, config)
-    weights = dict(tree_flatten(model.parameters()))
    save(
        args.mlx_path,
        model_path,
-        weights,
+        model,
        tokenizer,
        config,
-        hf_repo=args.model,
+        hf_repo=hf_repo,
    )
@@ -2,32 +2,38 @@

 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 tqdm import tqdm

-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.losses import kl_div_loss
+from mlx_lm.tuner.trainer import grad_checkpoint, 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,
+    save,
 )


+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,
@@ -35,8 +41,10 @@ def dwq_quantize(
    data,
    batch_size: int = 2,
    max_seq_length: int = 2048,
-    temperature: float = 0.5,
+    activation_layer_step: float = 0.25,
+    activation_loss_weight: float = 1.0,
    dtype: mx.Dtype = mx.bfloat16,
+    gradient_checkpoint: bool = False,
 ):
    group = mx.distributed.init()
    world_size = group.size()
@@ -49,22 +57,44 @@ def dwq_quantize(
    q_model.apply_to_modules(unfreeze)
    print_trainable_parameters(q_model)

-    def log_norm(x):
-        x = x * (1 / temperature)
-        return x - mx.logsumexp(x, axis=-1, keepdims=True)
+    layer_id_step = max(int(activation_layer_step * len(model.layers)), 1)
+    layer_ids = list(range(len(model.layers)))[layer_id_step::layer_id_step]

-    def loss_fn(params, x, targets, lengths):
+    for lid in layer_ids:
+        model.layers[lid] = Catcher(model.layers[lid])
+        q_model.layers[lid] = Catcher(q_model.layers[lid])
+
+    if gradient_checkpoint:
+        grad_checkpoint(q_model.layers[0])
+
+    def forward(model, inputs):
+        logits = model(inputs)
+        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 logits, extra_targets
+
+    def loss_fn(params, x, targets, extra_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:]
+        logits, q_extra_targets = forward(q_model, x)
+        losses = kl_div_loss(logits, targets)
+        mask = mx.arange(1, 1 + targets.shape[1]) < lengths[:, 1:]
        ntoks = mask.sum()
-        loss = (mask * losses).sum() / ntoks
+        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, lengths, params):
+    def step(inputs, targets, extra_targets, lengths, params):
        (loss, ntoks), grads = mx.value_and_grad(loss_fn)(
-            params, inputs, targets, lengths
+            params, inputs, targets, extra_targets, lengths
        )
        grads = nn.average_gradients(grads)
        params = opt.apply_gradients(grads, params)
@@ -76,53 +106,45 @@ def dwq_quantize(
        q_model.trainable_parameters(),
    )

-    avg_loss = None
+    total_loss = 0.0
+    total_tokens = 0
    tokens = 0
    tic = time.time()
-    for it, (batch, lengths) in enumerate(
-        iterate_batches(data, batch_size, max_seq_length)
+    for it, (batch, lengths) in (
+        pbar := tqdm(
+            enumerate(iterate_batches(data, batch_size, max_seq_length)),
+            total=len(data) // batch_size,
+        )
    ):
-        targets = log_norm(model(batch).astype(mx.float32))
-        mx.eval(targets)
-        loss, ntoks, params = step(batch, targets, lengths, params)
+        batch = batch[:, :-1]
+        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
+        total_loss += loss * ntoks
        if rank == 0:
-            print(
-                f"{it=}, {loss=:.3f}, {avg_loss=:.4f}, {tokens=}, {toks_per_sec=:.3f}",
-                flush=True,
-            )
+            pbar.set_description(desc=f"{loss=:.4f}")
+            if (it + 1) % 20 == 0:
+                toks_per_sec = tokens / (time.time() - tic)
+                peak_memory_gb = mx.get_peak_memory() / 1e9
+                avg_loss = total_loss / tokens
+                total_tokens += tokens
+                tqdm.write(
+                    f"{it=}, {avg_loss=:.4f}, {total_tokens=},"
+                    f" {toks_per_sec=:.3f}, {peak_memory_gb=:.3f}",
+                )
+                tic = time.time()
+                tokens = 0
+                total_loss = 0
    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):
+def load_data(tokenizer, data_path: str, num_samples: int, max_seq_length: int):
    args = types.SimpleNamespace(
        hf_dataset={
            "path": data_path,
@@ -134,12 +156,17 @@ def load_data(tokenizer, data_path: str, num_samples: int):
    )
    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 process(idx):
+        tokens, offset = dataset.process(dataset[idx])
+        return (tokens[:max_seq_length], offset)
+
+    return [process(i) for i in perm]


 def main():
    parser = argparse.ArgumentParser()
-    parser.add_argument("--model", "-m", default="Qwen/Qwen3-1.7B")
+    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."
@@ -156,13 +183,13 @@ def main():
    parser.add_argument(
        "--num-samples",
        type=int,
-        default=1024,
+        default=2048,
        help="Number of samples to use for training.",
    )
-    parser.add_argument("--max-seq-length", type=int, default=2048)
+    parser.add_argument("--max-seq-length", type=int, default=2049)
    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("--learning-rate", type=float, default=1e-6)
+    parser.add_argument("--batch-size", type=int, default=4)
    parser.add_argument(
        "--data-path",
        type=str,
@@ -170,10 +197,9 @@ def main():
        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.",
+        "--grad-checkpoint",
+        action="store_true",
+        help="Use gradient checkpointing to reduce memory use.",
    )
    args = parser.parse_args()

@@ -186,10 +212,12 @@ def main():
    np.random.seed(args.seed)
    mx.random.seed(args.seed)

-    model_path = get_model_path(args.model, revision=None)
+    model_path, hf_repo = 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)
+    calibration_data = load_data(
+        tokenizer, args.data_path, args.num_samples, args.max_seq_length
+    )

    if args.quantized_model is not None:
        q_model_path = get_model_path(args.quantized_model, revision=None)
@@ -211,6 +239,13 @@ def main():
        calibration_data,
        batch_size=args.batch_size,
        max_seq_length=args.max_seq_length,
-        temperature=args.temperature,
+        gradient_checkpoint=args.grad_checkpoint,
+    )
+    save(
+        args.mlx_path,
+        model_path,
+        q_model,
+        tokenizer,
+        config,
+        hf_repo=hf_repo,
    )
-    save_model(q_model, tokenizer, config, model_path, args.mlx_path, args.model)
@@ -0,0 +1,349 @@
+# Copyright © 2025 Apple Inc.
+
+import argparse
+import copy
+import json
+import math
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+from mlx.utils import tree_flatten, tree_map, tree_reduce, tree_unflatten
+from tqdm import tqdm
+
+from mlx_lm.quant.utils import load_data
+from mlx_lm.tuner.losses import kl_div_loss
+from mlx_lm.tuner.trainer import grad_checkpoint
+from mlx_lm.tuner.utils import get_total_parameters
+from mlx_lm.utils import (
+    compute_bits_per_weight,
+    fetch_from_hub,
+    get_model_path,
+    load,
+    quantize_model,
+    save,
+)
+
+
+def make_quant_predicate(config):
+    def quant_predicate(p, m, _):
+        if not hasattr(m, "to_quantized"):
+            return False
+        return config.get(p, True)
+
+    return quant_predicate
+
+
+def eval_ppl(model, data, batch_size=8):
+    all_loss = 0.0
+    ntoks = 0
+    for s in range(0, len(data), batch_size):
+        batch = data[s : s + batch_size]
+        logits = model(batch[:, :-1]).astype(mx.float32)
+        losses = nn.losses.cross_entropy(logits, batch[:, 1:])
+        all_loss += losses.sum().item()
+        ntoks += losses.size
+    ppl = math.exp(all_loss / ntoks)
+    return ppl
+
+
+def make_options(
+    low_bits, low_group_size, high_bits, high_group_size, include_bpw=True
+):
+    options = []
+    min_bpw = low_bits + 32 / low_group_size
+    max_bpw = high_bits + 32 / high_group_size
+    for b in range(low_bits, high_bits + 1):
+        for g in [32, 64, 128]:
+            cbpw = b + 32 / g
+            if b == 7 or not (min_bpw <= cbpw <= max_bpw):
+                continue
+            options.append({"bits": b, "group_size": g, "bpw": cbpw})
+    options.sort(key=lambda x: x["bpw"])
+    if not include_bpw:
+        for o in options:
+            o.pop("bpw")
+
+    return options
+
+
+def estimate_sensitivities(
+    model,
+    data,
+    low_bits,
+    low_group_size,
+    high_bits,
+    high_group_size,
+    batch_size: int = 4,
+    gradient_accum_dtype: mx.Dtype = mx.float32,
+    gradient_checkpoint: bool = False,
+):
+    def qdq(w, bits, group_size):
+        w, s, b = mx.quantize(w, bits=bits, group_size=group_size)
+        return mx.dequantize(w, scales=s, biases=b, bits=bits, group_size=group_size)
+
+    layers = tree_flatten(model.leaf_modules(), is_leaf=nn.Module.is_module)
+    layers = {k: l for k, l in layers if hasattr(l, "to_quantized")}
+    q_model = copy.deepcopy(model)
+    q_layers = copy.deepcopy(layers)
+    for l in q_layers.values():
+        l.weight = qdq(l.weight, low_bits, low_group_size)
+        # Freeze everything but the quantizable weight
+        l.freeze()
+        l.unfreeze(keys=["weight"])
+    q_model.freeze()
+    q_model.update_modules(tree_unflatten(list(q_layers.items())))
+
+    def loss_fn(batch, targets):
+        return kl_div_loss(q_model(batch), targets).mean()
+
+    if gradient_checkpoint:
+        grad_checkpoint(q_model.layers[0])
+
+    grad_accum = tree_map(
+        lambda x: mx.zeros(x.shape, dtype=gradient_accum_dtype),
+        q_model.trainable_parameters(),
+    )
+    for e, s in tqdm(
+        enumerate(range(0, len(data), batch_size)),
+        total=len(data) // batch_size,
+        desc="Estimating sensitivities",
+    ):
+        batch = data[s : s + batch_size]
+        targets = model(batch)
+        mx.eval(targets)
+        _, grads = nn.value_and_grad(q_model, loss_fn)(batch, targets)
+        grad_accum = tree_map(lambda x, y: x + y, grad_accum, grads)
+        del grads
+        mx.eval(grad_accum)
+
+    options = make_options(low_bits, low_group_size, high_bits, high_group_size)
+    current_bpw = options[0]["bpw"]
+
+    def compute_sensitivity(gradient, low_q_weight, original_weight):
+        n_batches = (len(data) + batch_size - 1) // batch_size
+        gradient = gradient / n_batches
+        scores = [{"loss_change": 0, "extra_bits": 0}]
+        for opt in options[1:]:
+            extra_bits = (opt["bpw"] - current_bpw) * original_weight.size
+            other_weight = qdq(original_weight, opt["bits"], opt["group_size"])
+            loss_change = (gradient * (low_q_weight - other_weight)).sum()
+            scores.append({"loss_change": loss_change, "extra_bits": extra_bits})
+        return scores
+
+    sensitivities = tree_map(
+        compute_sensitivity,
+        grad_accum,
+        q_model.parameters(),
+        model.parameters(),
+    )
+    mx.eval(sensitivities)
+
+    sensitivities = [
+        (k.replace(".weight", ""), s.item() if isinstance(s, mx.array) else s)
+        for k, s in tree_flatten(sensitivities)
+    ]
+
+    return sensitivities
+
+
+def compute_bit_budget(model, target_bpw):
+    model_bytes = tree_reduce(
+        lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, model, 0
+    )
+    model_params = get_total_parameters(model)
+
+    return model_params * target_bpw - model_bytes * 8
+
+
+def estimate_threshold(
+    model,
+    sensitivities,
+    target_bpw,
+    low_bits,
+    low_group_size,
+    high_bits,
+    high_group_size,
+):
+    options = make_options(
+        low_bits, low_group_size, high_bits, high_group_size, include_bpw=False
+    )
+    sensitivities = tree_flatten(
+        tree_unflatten(list(sensitivities.items())),
+        is_leaf=lambda x: isinstance(x, list) and "loss_change" in x[0],
+    )
+
+    q_model = copy.deepcopy(model)
+    nn.quantize(q_model, group_size=low_group_size, bits=low_bits)
+    budget = int(compute_bit_budget(q_model, target_bpw))
+    benefit_map = {}
+
+    def benefit(layer, option, budget):
+        if (layer, option, budget) in benefit_map:
+            return benefit_map[layer, option, budget]
+
+        stack = [(layer, option, budget)]
+        while stack:
+            layer, option, budget = stack[-1]
+
+            if budget <= 0 or layer < 0 or option < 0:
+                benefit_map[layer, option, budget] = 0
+                stack.pop()
+                continue
+
+            # We either not use this option
+            prev_layer = layer if option > 0 else layer - 1
+            prev_option = (option if option > 0 else len(options)) - 1
+            if (prev_layer, prev_option, budget) not in benefit_map:
+                stack.append((prev_layer, prev_option, budget))
+                continue
+            a = benefit_map[prev_layer, prev_option, budget]
+
+            # Or we use it so we have less budget for before
+            b = float("-inf")
+            info = sensitivities[layer][1][option]
+            prev_layer = layer - 1
+            prev_option = len(options) - 1
+            prev_budget = budget - info["extra_bits"]
+            if (
+                prev_layer,
+                prev_option,
+                prev_budget,
+            ) not in benefit_map and prev_budget >= 0:
+                stack.append((prev_layer, prev_option, prev_budget))
+                continue
+            if prev_budget >= 0:
+                b = benefit_map[prev_layer, prev_option, prev_budget]
+                b += info["loss_change"]
+
+            benefit_map[layer, option, budget] = max(a, b)
+            stack.pop()
+
+        return benefit_map[layer, option, budget]
+
+    def backtrack(layer, budget):
+        selected = []
+        while layer >= 0:
+            prev_benefit = benefit(layer - 1, len(options) - 1, budget)
+            option_benefits = [benefit(layer, i, budget) for i in range(len(options))]
+            idx, v = max(enumerate(option_benefits), key=lambda x: x[1] - prev_benefit)
+            info = sensitivities[layer][1][idx]
+            if v != 0:
+                budget -= info["extra_bits"]
+                selected.append((layer, idx))
+            layer -= 1
+        return selected[::-1]
+
+    selected = backtrack(len(sensitivities) - 1, budget)
+    config = {sensitivities[l][0]: options[i] for l, i in selected}
+
+    return config
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", "-m", default="Qwen/Qwen3-0.6B-base")
+    parser.add_argument(
+        "--mlx-path", default="mlx_model", help="Path to save the model"
+    )
+    parser.add_argument("--seed", type=int, default=123)
+    parser.add_argument(
+        "--sensitivities",
+        type=str,
+        default=None,
+        help="Path to a pre-computed sensitivity JSON file.",
+    )
+    parser.add_argument(
+        "--target-bpw", type=float, default=5.0, help="Target bits per weight."
+    )
+    parser.add_argument("--low-bits", type=int, default=4)
+    parser.add_argument("--low-group-size", type=int, default=128)
+    parser.add_argument("--high-bits", type=int, default=5)
+    parser.add_argument("--high-group-size", type=int, default=32)
+    parser.add_argument(
+        "--report-ppl",
+        action="store_true",
+        help="Compute the perplexity of the base and quantized models.",
+    )
+    parser.add_argument(
+        "--grad-checkpoint",
+        action="store_true",
+        help="Use gradient checkpointing to reduce memory use.",
+    )
+    parser.add_argument(
+        "--accumulation-dtype",
+        default="float32",
+        choices=["float32", "bfloat16"],
+        help="What type to use to accumulate the gradients for the sensitivities",
+    )
+    args = parser.parse_args()
+
+    group = mx.distributed.init()
+
+    if args.sensitivities is None:
+        model, tokenizer = load(args.model)
+        mx.random.seed(args.seed)
+        data = load_data(tokenizer, num_samples=-1, sequence_length=512)
+
+        sensitivities = estimate_sensitivities(
+            model,
+            data,
+            args.low_bits,
+            args.low_group_size,
+            args.high_bits,
+            args.high_group_size,
+            gradient_accum_dtype=getattr(mx, args.accumulation_dtype),
+            gradient_checkpoint=args.grad_checkpoint,
+        )
+        model_name = args.model.replace("/", "_")
+        with open(f"{model_name}_sensitivities.json", "w") as fid:
+            json.dump(sensitivities, fid)
+    else:
+        with open(args.sensitivities, "r") as fid:
+            sensitivities = json.load(fid)
+
+    sensitivities = dict(sensitivities)
+    model_path, hf_repo = get_model_path(args.model, revision=None)
+    model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
+    mx.random.seed(args.seed)
+    data = load_data(tokenizer, num_samples=-1, sequence_length=512)
+
+    if args.report_ppl:
+        ppl = eval_ppl(model, data)
+        print(f"Original PPL: {ppl:.3f}")
+
+    quant_config = estimate_threshold(
+        model,
+        sensitivities,
+        target_bpw=args.target_bpw,
+        low_bits=args.low_bits,
+        low_group_size=args.low_group_size,
+        high_bits=args.high_bits,
+        high_group_size=args.high_group_size,
+    )
+
+    model, config = quantize_model(
+        model,
+        config,
+        q_group_size=args.low_group_size,
+        q_bits=args.low_bits,
+        quant_predicate=make_quant_predicate(quant_config),
+    )
+
+    if args.report_ppl:
+        ppl = eval_ppl(model, data)
+        print(f"Quantized PPL: {ppl:.3f}")
+
+    save(
+        args.mlx_path,
+        model_path,
+        model,
+        tokenizer,
+        config,
+        hf_repo=hf_repo,
+    )
+    print(f"Peak memory used: {mx.get_peak_memory() / 1000**3:.3f}GB")
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,26 @@
+# Copyright © 2025 Apple Inc.
+
+from pathlib import Path
+
+import mlx.core as mx
+
+
+def load_data(tokenizer, num_samples: int, sequence_length: int) -> mx.array:
+    save_dir = Path.home() / ".cache/mlx-lm/calibration_v5.txt"
+    if not save_dir.exists():
+        from urllib import request
+
+        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])
+    if num_samples > 0:
+        segments = segments[:num_samples]
+    return tokens[segments]
@@ -184,8 +184,12 @@ def apply_min_p(
    selected_logprobs = mx.where(tokens_to_remove, -float("inf"), sorted_logprobs)

    # Create a mapping to rearrange back to original indices
-    # Use argsort of sorted_indices to get the inverse permutation
-    inverse_indices = mx.argsort(sorted_indices, axis=-1)
+    inverse_indices = mx.put_along_axis(
+        mx.zeros_like(sorted_indices),
+        sorted_indices,
+        mx.arange(sorted_indices.shape[-1], dtype=sorted_indices.dtype),
+        axis=-1,
+    )

    # Rearrange selected_logprobs back to original order
    original_order_logprobs = mx.take_along_axis(
@@ -196,40 +200,39 @@ def apply_min_p(


@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
-def apply_top_p(logits: mx.array, top_p: float) -> mx.array:
+def apply_top_p(logprobs: mx.array, top_p: float) -> mx.array:
    """
    Apply top-p (nucleus) sampling to logits.

    Args:
-        logits: The logits from the model's output.
+        logprobs: A vector of log probabilities.
        top_p: The cumulative probability threshold for top-p filtering.
    Returns:
        token selected based on the top-p criterion.
    """
    # 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)
+    probs = mx.exp(logprobs)
+    # sort in ascending order
+    sorted_indices = mx.argsort(logprobs, axis=-1)
    sorted_probs = mx.take_along_axis(probs, sorted_indices, axis=-1)

    cumulative_probs = mx.cumsum(sorted_probs, axis=-1)

-    # select tokens with cumulative probs below threshold
-    top_probs = mx.where(
-        cumulative_probs > 1 - top_p,
-        sorted_probs,
-        0,
+    # Rearrange cumulative probs back to original order
+    inverse_indices = mx.put_along_axis(
+        mx.zeros_like(sorted_indices),
+        sorted_indices,
+        mx.arange(sorted_indices.shape[-1], dtype=sorted_indices.dtype),
+        axis=-1,
    )
+    cumulative_probs = mx.take_along_axis(cumulative_probs, inverse_indices, axis=-1)

-    # Create a mapping to rearrange back to original indices
-    # Use argsort of sorted_indices to get the inverse permutation
-    inverse_indices = mx.argsort(sorted_indices, axis=-1)
-
-    # Rearrange top_probs back to original order
-    original_order_probs = mx.take_along_axis(top_probs, inverse_indices, axis=-1)
-
-    # Convert back to logits and return
-    return mx.log(original_order_probs)
+    # select tokens with cumulative probs below threshold
+    return mx.where(
+        cumulative_probs > 1 - top_p,
+        logprobs,
+        -float("inf"),
+    )


@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
@@ -141,6 +141,8 @@ def process_message_content(messages):
            if len(text_fragments) != len(content):
                raise ValueError("Only 'text' content type is supported.")
            message["content"] = "".join(text_fragments)
+        elif content is None:
+            message["content"] = ""


@dataclass
@@ -307,13 +309,21 @@ class APIHandler(BaseHTTPRequestHandler):
        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.num_draft_tokens = self.body.get(
+            "num_draft_tokens", self.model_provider.cli_args.num_draft_tokens
+        )
        self.adapter = self.body.get("adapters", None)
        self.max_tokens = self.body.get("max_completion_tokens", None)
        if self.max_tokens is None:
-            self.max_tokens = self.body.get("max_tokens", 512)
-        self.temperature = self.body.get("temperature", 0.0)
-        self.top_p = self.body.get("top_p", 1.0)
+            self.max_tokens = self.body.get(
+                "max_tokens", self.model_provider.cli_args.max_tokens
+            )
+        self.temperature = self.body.get(
+            "temperature", self.model_provider.cli_args.temp
+        )
+        self.top_p = self.body.get("top_p", self.model_provider.cli_args.top_p)
+        self.top_k = self.body.get("top_k", self.model_provider.cli_args.top_k)
+        self.min_p = self.body.get("min_p", self.model_provider.cli_args.min_p)
        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)
@@ -370,6 +380,15 @@ class APIHandler(BaseHTTPRequestHandler):
        if not isinstance(self.top_p, (float, int)) or self.top_p < 0 or self.top_p > 1:
            raise ValueError("top_p must be a float between 0 and 1")

+        if not isinstance(self.top_k, int) or self.top_k < 0:
+            raise ValueError("top_k must be a non-negative integer")
+
+        if not isinstance(self.min_p, (float, int)) or self.min_p < 0 or self.min_p > 1:
+            raise ValueError("min_p must be a float between 0 and 1")
+
+        if not isinstance(self.num_draft_tokens, int) or self.num_draft_tokens < 0:
+            raise ValueError("num_draft_tokens must be a non-negative integer")
+
        if (
            not isinstance(self.repetition_penalty, (float, int))
            or self.repetition_penalty < 0
@@ -418,6 +437,7 @@ class APIHandler(BaseHTTPRequestHandler):
        token_logprobs: Optional[List[float]] = None,
        top_tokens: Optional[List[Dict[int, float]]] = None,
        tokens: Optional[List[int]] = None,
+        tool_calls: Optional[List[str]] = None,
    ) -> dict:
        """
        Generate a single response packet based on response type (stream or
@@ -436,13 +456,26 @@ class APIHandler(BaseHTTPRequestHandler):
            top_tokens (Optional[List[Dict[int, float]]]): List of dictionaries mapping
              tokens to logprobs for the top N tokens at each token position.
            tokens (Optional[List[int]]): List of tokens to return with logprobs structure
+            tool_calls (Optional[List[str]]): List of tool calls.

        Returns:
            dict: A dictionary containing the response, in the same format as
              OpenAI's API.
        """
-        token_logprobs = token_logprobs if token_logprobs else []
-        top_logprobs = top_tokens if top_tokens else []
+        token_logprobs = token_logprobs or []
+        top_logprobs = top_tokens or []
+        tool_calls = tool_calls or []
+
+        def parse_function(tool_text):
+            tool_call = json.loads(tool_text.strip())
+            return {
+                "function": {
+                    "name": tool_call.get("name", None),
+                    "arguments": json.dumps(tool_call.get("arguments", "")),
+                },
+                "type": "function",
+                "id": None,
+            }

        # Static response
        response = {
@@ -460,7 +493,7 @@ class APIHandler(BaseHTTPRequestHandler):
                        "tokens": tokens,
                    },
                    "finish_reason": finish_reason,
-                }
+                },
            ],
        }

@@ -484,7 +517,11 @@ class APIHandler(BaseHTTPRequestHandler):
        # Add dynamic response
        if self.object_type.startswith("chat.completion"):
            key_name = "delta" if self.stream else "message"
-            choice[key_name] = {"role": "assistant", "content": text}
+            choice[key_name] = {
+                "role": "assistant",
+                "content": text,
+                "tool_calls": [parse_function(tool_text) for tool_text in tool_calls],
+            }
        elif self.object_type == "text_completion":
            choice.update(text=text)
        else:
@@ -530,6 +567,9 @@ class APIHandler(BaseHTTPRequestHandler):
        prompt_len = len(prompt)
        com_prefix_len = common_prefix_len(self.prompt_cache.tokens, prompt)

+        # Leave at least one token in the prompt
+        com_prefix_len = min(com_prefix_len, len(prompt) - 1)
+
        # Condition 1: Model changed or no common prefix at all. Reset cache.
        if (
            self.prompt_cache.model_key != self.model_provider.model_key
@@ -603,6 +643,8 @@ class APIHandler(BaseHTTPRequestHandler):
        sampler = make_sampler(
            self.temperature,
            top_p=self.top_p,
+            top_k=self.top_k,
+            min_p=self.min_p,
            xtc_probability=self.xtc_probability,
            xtc_threshold=self.xtc_threshold,
            xtc_special_tokens=[
@@ -616,6 +658,10 @@ class APIHandler(BaseHTTPRequestHandler):
            self.repetition_context_size,
        )

+        tool_calls = []
+        tool_text = ""
+        in_tool_call = False
+        segment = ""
        for gen_response in stream_generate(
            model=self.model,
            tokenizer=self.tokenizer,
@@ -627,9 +673,23 @@ class APIHandler(BaseHTTPRequestHandler):
            draft_model=self.model_provider.draft_model,
            num_draft_tokens=self.num_draft_tokens,
        ):
-            segment = gen_response.text
-            text += segment
-            logging.debug(text)
+            logging.debug(gen_response.text)
+
+            if (
+                self.tokenizer.has_tool_calling
+                and gen_response.text == self.tokenizer.tool_call_start
+            ):
+                in_tool_call = True
+            elif in_tool_call:
+                if gen_response.text == self.tokenizer.tool_call_end:
+                    tool_calls.append(tool_text)
+                    tool_text = ""
+                    in_tool_call = False
+                else:
+                    tool_text += gen_response.text
+            else:
+                text += gen_response.text
+                segment += gen_response.text
            token = gen_response.token
            logprobs = gen_response.logprobs
            tokens.append(token)
@@ -653,9 +713,10 @@ class APIHandler(BaseHTTPRequestHandler):
                        tokens[-stop_condition.trim_length :]
                    )
                    text = text[: -len(stop_sequence_suffix)]
+                segment = ""
                break

-            if self.stream:
+            if self.stream and not in_tool_call:
                # If the end of tokens overlaps with a stop sequence, generate new
                # tokens until we know if the stop sequence is hit or not
                if any(
@@ -665,10 +726,14 @@ class APIHandler(BaseHTTPRequestHandler):
                    )
                ):
                    continue
-                elif segment:
-                    response = self.generate_response(segment, None)
+                elif segment or tool_calls:
+                    response = self.generate_response(
+                        segment, None, tool_calls=tool_calls
+                    )
                    self.wfile.write(f"data: {json.dumps(response)}\n\n".encode())
                    self.wfile.flush()
+                    segment = ""
+                    tool_calls = []

        self.prompt_cache.tokens.extend(tokens)

@@ -677,7 +742,9 @@ class APIHandler(BaseHTTPRequestHandler):
        logging.debug(f"Peak memory: {gen_response.peak_memory:.3f} GB")

        if self.stream:
-            response = self.generate_response(segment, finish_reason)
+            response = self.generate_response(
+                segment, finish_reason, tool_calls=tool_calls
+            )
            self.wfile.write(f"data: {json.dumps(response)}\n\n".encode())
            self.wfile.flush()
            if self.stream_options is not None and self.stream_options["include_usage"]:
@@ -695,6 +762,7 @@ class APIHandler(BaseHTTPRequestHandler):
                token_logprobs=token_logprobs,
                top_tokens=top_tokens,
                tokens=tokens,
+                tool_calls=tool_calls,
            )
            response_json = json.dumps(response).encode()
            indent = "\t"  # Backslashes can't be inside of f-strings
@@ -744,8 +812,9 @@ class APIHandler(BaseHTTPRequestHandler):
            process_message_content(messages)
            prompt = self.tokenizer.apply_chat_template(
                messages,
-                body.get("tools", None),
+                body.get("tools") or None,
                add_generation_prompt=True,
+                **self.model_provider.cli_args.chat_template_args,
            )
        else:
            prompt = convert_chat(body["messages"], body.get("role_mapping"))
@@ -772,11 +841,23 @@ class APIHandler(BaseHTTPRequestHandler):
        """
        if self.path == "/v1/models":
            self.handle_models_request()
+        elif self.path == "/health":
+            self.handle_health_check()
        else:
            self._set_completion_headers(404)
            self.end_headers()
            self.wfile.write(b"Not Found")

+    def handle_health_check(self):
+        """
+        Handle a GET request for the /health endpoint.
+        """
+        self._set_completion_headers(200)
+        self.end_headers()
+
+        self.wfile.write('{"status": "ok"}'.encode())
+        self.wfile.flush()
+
    def handle_models_request(self):
        """
        Handle a GET request for the /v1/models endpoint.
@@ -878,6 +959,12 @@ def main():
        help="A model to be used for speculative decoding.",
        default=None,
    )
+    parser.add_argument(
+        "--num-draft-tokens",
+        type=int,
+        help="Number of tokens to draft when using speculative decoding.",
+        default=3,
+    )
    parser.add_argument(
        "--trust-remote-code",
        action="store_true",
@@ -902,6 +989,42 @@ def main():
        action="store_true",
        help="Use the default chat template",
    )
+    parser.add_argument(
+        "--temp",
+        type=float,
+        default=0.0,
+        help="Default sampling temperature (default: 0.0)",
+    )
+    parser.add_argument(
+        "--top-p",
+        type=float,
+        default=1.0,
+        help="Default nucleus sampling top-p (default: 1.0)",
+    )
+    parser.add_argument(
+        "--top-k",
+        type=int,
+        default=0,
+        help="Default top-k sampling (default: 0, disables top-k)",
+    )
+    parser.add_argument(
+        "--min-p",
+        type=float,
+        default=0.0,
+        help="Default min-p sampling (default: 0.0, disables min-p)",
+    )
+    parser.add_argument(
+        "--max-tokens",
+        type=int,
+        default=512,
+        help="Default maximum number of tokens to generate (default: 512)",
+    )
+    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()

    logging.basicConfig(
@@ -3,7 +3,7 @@ from functools import partial
 from json import JSONDecodeError
 from typing import List

-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, PreTrainedTokenizerFast


 class StreamingDetokenizer:
@@ -91,6 +91,7 @@ class NaiveStreamingDetokenizer(StreamingDetokenizer):
            self._current_text = self._tokenizer.decode(self._current_tokens)
            if (
                self._tokenizer.clean_up_tokenization_spaces
+                and len(self._current_text) > 0
                and self._current_text[-1] == " "
            ):
                self._current_text = self._current_text[:-1]
@@ -267,6 +268,28 @@ class TokenizerWrapper:
            if eos_token_ids is not None
            else {tokenizer.eos_token_id}
        )
+        self._think_start = None
+        self._think_end = None
+        self._tool_call_start = None
+        self._tool_call_end = None
+
+        THINK_TOKENS = [("<think>", "</think>")]
+        TOOL_CALL_TOKENS = [("<tool_call>", "</tool_call>")]
+
+        vocab = tokenizer.get_vocab()
+        for think_start, think_end in THINK_TOKENS:
+            if think_start in vocab and think_end in vocab:
+                self._think_start = think_start
+                self._think_end = think_end
+                break
+        if tokenizer.chat_template and '"tool"' in tokenizer.chat_template:
+            self._tool_call_start = ""
+            self._tool_call_end = ""
+            for tool_call_start, tool_call_end in TOOL_CALL_TOKENS:
+                if tool_call_start in vocab and tool_call_end in vocab:
+                    self._tool_call_start = tool_call_start
+                    self._tool_call_end = tool_call_end
+                    break

    def add_eos_token(self, token: str):
        token_id = None
@@ -280,6 +303,30 @@ class TokenizerWrapper:

        self._eos_token_ids.add(token_id)

+    @property
+    def has_thinking(self):
+        return self._think_start is not None
+
+    @property
+    def think_start(self):
+        return self._think_start
+
+    @property
+    def think_end(self):
+        return self._think_end
+
+    @property
+    def has_tool_calling(self):
+        return self._tool_call_start is not None
+
+    @property
+    def tool_call_start(self):
+        return self._tool_call_start
+
+    @property
+    def tool_call_end(self):
+        return self._tool_call_end
+
    def __getattr__(self, attr):
        if attr == "detokenizer":
            return self._detokenizer
@@ -302,6 +349,35 @@ class TokenizerWrapper:
            setattr(self._tokenizer, attr, value)


+class NewlineTokenizer(PreTrainedTokenizerFast):
+    """A tokenizer that replaces newlines with <n> and <n> with new line."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def _preprocess_text(self, text):
+        return text.replace("\n", "<n>")
+
+    def _postprocess_text(self, text):
+        return text.replace("<n>", "\n")
+
+    def encode(self, text, **kwargs):
+        return super().encode(self._preprocess_text(text), **kwargs)
+
+    def encode_batch(self, texts, **kwargs):
+        return super().encode_batch([self._preprocess_text(t) for t in texts], **kwargs)
+
+    def decode(self, *args, **kwargs):
+        return self._postprocess_text(super().decode(*args, **kwargs))
+
+    def batch_decode(self, *args, **kwargs):
+        decoded = super().batch_decode(*args, **kwargs)
+        return [self._postprocess_text(d) for d in decoded]
+
+
+AutoTokenizer.register("NewlineTokenizer", fast_tokenizer_class=NewlineTokenizer)
+
+
 def _match(a, b):
    if type(a) != type(b):
        return False
@@ -0,0 +1,43 @@
+# Copyright © 2024 Apple Inc.
+
+try:
+    import wandb
+except ImportError:
+    wandb = None
+
+
+class TrainingCallback:
+
+    def on_train_loss_report(self, train_info: dict):
+        """Called to report training loss at specified intervals."""
+        pass
+
+    def on_val_loss_report(self, val_info: dict):
+        """Called to report validation loss at specified intervals or the beginning."""
+        pass
+
+
+class WandBCallback(TrainingCallback):
+    def __init__(
+        self,
+        project_name: str,
+        log_dir: str,
+        config: dict,
+        wrapped_callback: TrainingCallback = None,
+    ):
+        if wandb is None:
+            raise ImportError(
+                "wandb is not installed. Please install it to use WandBCallback."
+            )
+        self.wrapped_callback = wrapped_callback
+        wandb.init(project=project_name, dir=log_dir, config=config)
+
+    def on_train_loss_report(self, train_info: dict):
+        wandb.log(train_info, step=train_info.get("iteration"))
+        if self.wrapped_callback:
+            self.wrapped_callback.on_train_loss_report(train_info)
+
+    def on_val_loss_report(self, val_info: dict):
+        wandb.log(val_info, step=val_info.get("iteration"))
+        if self.wrapped_callback:
+            self.wrapped_callback.on_val_loss_report(val_info)
@@ -1,7 +1,9 @@
+# Copyright © 2024 Apple Inc.
+
 import json
 import types
 from pathlib import Path
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List

 from transformers import PreTrainedTokenizer

@@ -25,7 +27,7 @@ class TextDataset:
        d = self.tokenizer.encode(d[self.text_key])
        if d[-1] != self.tokenizer.eos_token_id:
            d.append(self.tokenizer.eos_token_id)
-        return d
+        return (d, 0)

    def __getitem__(self, idx: int):
        return self._data[idx]
@@ -61,7 +63,7 @@ class ChatDataset:
            offset = len(self.tokenizer.apply_chat_template(messages, tools=tools))
            return (tokens, offset)
        else:
-            return tokens
+            return (tokens, 0)

    def __getitem__(self, idx: int):
        return self._data[idx]
@@ -106,7 +108,7 @@ class CompletionsDataset:
            )
            return (tokens, offset)

-        return tokens
+        return (tokens, 0)

    def __getitem__(self, idx: int):
        return self._data[idx]
@@ -180,7 +182,7 @@ def create_dataset(
    else:
        raise ValueError(
            "Unsupported data format, check the supported formats here:\n"
-            "https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/LORA.md#data."
+            "https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/LORA.md#Data."
        )


@@ -20,7 +20,7 @@ class LoRALinear(nn.Module):
        # on linear and quantized linear
        output_dims, input_dims = linear.weight.shape
        if isinstance(linear, nn.QuantizedLinear):
-            input_dims *= 32 // linear.bits
+            input_dims = input_dims * 32 // linear.bits
        lora_lin = LoRALinear(
            input_dims=input_dims,
            output_dims=output_dims,
@@ -52,9 +52,8 @@ class LoRALinear(nn.Module):
        output_dims, input_dims = weight.shape
        fused_linear = nn.Linear(input_dims, output_dims, bias=bias)

-        lora_b = (self.scale * self.lora_b.T).astype(dtype)
-        lora_a = self.lora_a.T.astype(dtype)
-        fused_linear.weight = weight + lora_b @ lora_a
+        delta = ((self.scale * self.lora_b.T) @ self.lora_a.T).astype(dtype)
+        fused_linear.weight = weight + delta
        if bias:
            fused_linear.bias = linear.bias

@@ -203,7 +202,7 @@ class LoRAEmbedding(nn.Module):
    ):
        num_embeddings, dims = embedding.weight.shape
        if isinstance(embedding, nn.QuantizedEmbedding):
-            dims *= 32 // embedding.bits
+            dims = dims * 32 // embedding.bits
        lora_embedding = LoRAEmbedding(
            num_embeddings=num_embeddings,
            dims=dims,
@@ -0,0 +1,378 @@
+# Copyright © 2025 Apple Inc.
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+def _make_kl_forward_kernel():
+    source = """
+    constexpr int M = 4;
+    constexpr int block = 1024 * M;
+    constexpr int full_blocks = V / block;
+    constexpr int extra = V - full_blocks * block;
+
+    threadgroup float shared[32 * 2];
+
+    uint out_idx = threadgroup_position_in_grid.y;
+    uint simd_lane_id = thread_index_in_simdgroup;
+    uint simd_group_id = simdgroup_index_in_threadgroup;
+
+    logits_q += out_idx * V;
+    logits_p += out_idx * V;
+    out += out_idx;
+
+    float lse_q_minus_p;
+    float lse_p;
+
+    {
+        float max_q = -1e30;
+        float max_p = -1e30;
+        float sum_exp_q = 0;
+        float sum_exp_p = 0;
+
+        int offset = thread_index_in_threadgroup * M;
+        for (int i = 0; i < full_blocks; i++) {
+            // Read and update q and p
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = logits_q[offset + j];
+                vals_p[j] = logits_p[offset + j];
+            }
+            float prev_max_q = max_q;
+            float prev_max_p = max_p;
+            for (int j=0; j<M; j++) {
+                max_q = max(max_q, vals_q[j]);
+                max_p = max(max_p, vals_p[j]);
+            }
+            sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+            sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+            for (int j=0; j<M; j++) {
+                sum_exp_q += metal::fast::exp(vals_q[j] - max_q);
+                sum_exp_p += metal::fast::exp(vals_p[j] - max_p);
+            }
+
+            // Move to the next block
+            offset += block;
+        }
+        if (extra > 0) {
+            // Read and update q and p
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j < M; j++) {
+                vals_q[j] = (offset + j < V) ? logits_q[offset + j] : -1e30;
+                vals_p[j] = (offset + j < V) ? logits_p[offset + j] : -1e30;
+            }
+            float prev_max_q = max_q;
+            float prev_max_p = max_p;
+            for (int j=0; j<M; j++) {
+                max_q = max(max_q, vals_q[j]);
+                max_p = max(max_p, vals_p[j]);
+            }
+            sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+            sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+            for (int j=0; j<M; j++) {
+                sum_exp_q += metal::fast::exp(vals_q[j] - max_q);
+                sum_exp_p += metal::fast::exp(vals_p[j] - max_p);
+            }
+        }
+
+        // Share the maxs across the threadgroup
+        float prev_max_q = max_q;
+        float prev_max_p = max_p;
+        max_q = simd_max(max_q);
+        max_p = simd_max(max_p);
+        if (simd_lane_id == 0) {
+            shared[simd_group_id * 2 + 0] = max_q;
+            shared[simd_group_id * 2 + 1] = max_p;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        max_q = shared[simd_lane_id * 2 + 0];
+        max_p = shared[simd_lane_id * 2 + 1];
+        max_q = simd_max(max_q);
+        max_p = simd_max(max_p);
+
+        // Share the sum_exp across the threadgroup
+        sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+        sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+        sum_exp_q = simd_sum(sum_exp_q);
+        sum_exp_p = simd_sum(sum_exp_p);
+        if (simd_lane_id == 0) {
+            shared[simd_group_id * 2 + 0] = sum_exp_q;
+            shared[simd_group_id * 2 + 1] = sum_exp_p;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        sum_exp_q = shared[simd_lane_id * 2 + 0];
+        sum_exp_p = shared[simd_lane_id * 2 + 1];
+        sum_exp_q = simd_sum(sum_exp_q);
+        sum_exp_p = simd_sum(sum_exp_p);
+
+        lse_p = max_p + metal::fast::log(sum_exp_p);
+        lse_q_minus_p = max_q + metal::fast::log(sum_exp_q) - lse_p;
+    }
+
+    threadgroup_barrier(mem_flags::mem_none);
+
+    {
+        float kl = 0;
+
+        int offset = thread_index_in_threadgroup * M;
+        for (int i = 0; i < full_blocks; i++) {
+            // Read and add to the kl
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = logits_q[offset + j];
+                vals_p[j] = logits_p[offset + j];
+            }
+
+            for (int j=0; j<M; j++) {
+                kl += metal::fast::exp(vals_p[j] - lse_p) * (vals_p[j] - vals_q[j] + lse_q_minus_p);
+            }
+
+            // Move to the next block
+            offset += block;
+        }
+        if (extra > 0) {
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = (offset + j < V) ? logits_q[offset + j] : -1e30;
+                vals_p[j] = (offset + j < V) ? logits_p[offset + j] : -1e30;
+            }
+
+            for (int j=0; j<M; j++) {
+                kl += metal::fast::exp(vals_p[j] - lse_p) * (vals_p[j] - vals_q[j] + lse_q_minus_p);
+            }
+        }
+
+        // Add the kl across the threadgroup
+        kl = simd_sum(kl);
+        if (simd_lane_id == 0) {
+            shared[simd_group_id] = kl;
+        }
+        threadgroup_barrier(mem_flags::mem_none);
+        kl = shared[simd_lane_id];
+        kl = simd_sum(kl);
+
+        if (thread_index_in_threadgroup == 0) {
+            out[0] = static_cast<T>(kl);
+        }
+    }
+    """
+
+    return mx.fast.metal_kernel(
+        name="kl_forward",
+        input_names=["logits_q", "logits_p"],
+        output_names=["out"],
+        source=source,
+        ensure_row_contiguous=True,
+    )
+
+
+def _make_kl_backward_kernel():
+    source = """
+    constexpr int M = 4;
+    constexpr int block = 1024 * M;
+    constexpr int full_blocks = V / block;
+    constexpr int extra = V - full_blocks * block;
+
+    threadgroup float shared[32 * 2];
+
+    uint out_idx = threadgroup_position_in_grid.y;
+    uint simd_lane_id = thread_index_in_simdgroup;
+    uint simd_group_id = simdgroup_index_in_threadgroup;
+
+    logits_q += out_idx * V;
+    logits_p += out_idx * V;
+    out += out_idx * V;
+    cotan += out_idx;
+
+    float lse_q;
+    float lse_p;
+
+    {
+        float max_q = -1e30;
+        float max_p = -1e30;
+        float sum_exp_q = 0;
+        float sum_exp_p = 0;
+
+        int offset = thread_index_in_threadgroup * M;
+        for (int i = 0; i < full_blocks; i++) {
+            // Read and update q and p
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = logits_q[offset + j];
+                vals_p[j] = logits_p[offset + j];
+            }
+            float prev_max_q = max_q;
+            float prev_max_p = max_p;
+            for (int j=0; j<M; j++) {
+                max_q = max(max_q, vals_q[j]);
+                max_p = max(max_p, vals_p[j]);
+            }
+            sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+            sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+            for (int j=0; j<M; j++) {
+                sum_exp_q += metal::fast::exp(vals_q[j] - max_q);
+                sum_exp_p += metal::fast::exp(vals_p[j] - max_p);
+            }
+
+            // Move to the next block
+            offset += block;
+        }
+        if (extra > 0) {
+            // Read and update q and p
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j < M; j++) {
+                vals_q[j] = (offset + j < V) ? logits_q[offset + j] : -1e30;
+                vals_p[j] = (offset + j < V) ? logits_p[offset + j] : -1e30;
+            }
+            float prev_max_q = max_q;
+            float prev_max_p = max_p;
+            for (int j=0; j<M; j++) {
+                max_q = max(max_q, vals_q[j]);
+                max_p = max(max_p, vals_p[j]);
+            }
+            sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+            sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+            for (int j=0; j<M; j++) {
+                sum_exp_q += metal::fast::exp(vals_q[j] - max_q);
+                sum_exp_p += metal::fast::exp(vals_p[j] - max_p);
+            }
+        }
+
+        // Share the maxs across the threadgroup
+        float prev_max_q = max_q;
+        float prev_max_p = max_p;
+        max_q = simd_max(max_q);
+        max_p = simd_max(max_p);
+        if (simd_lane_id == 0) {
+            shared[simd_group_id * 2 + 0] = max_q;
+            shared[simd_group_id * 2 + 1] = max_p;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        max_q = shared[simd_lane_id * 2 + 0];
+        max_p = shared[simd_lane_id * 2 + 1];
+        max_q = simd_max(max_q);
+        max_p = simd_max(max_p);
+
+        // Share the sum_exp across the threadgroup
+        sum_exp_q *= metal::fast::exp(prev_max_q - max_q);
+        sum_exp_p *= metal::fast::exp(prev_max_p - max_p);
+        sum_exp_q = simd_sum(sum_exp_q);
+        sum_exp_p = simd_sum(sum_exp_p);
+        if (simd_lane_id == 0) {
+            shared[simd_group_id * 2 + 0] = sum_exp_q;
+            shared[simd_group_id * 2 + 1] = sum_exp_p;
+        }
+        threadgroup_barrier(mem_flags::mem_threadgroup);
+        sum_exp_q = shared[simd_lane_id * 2 + 0];
+        sum_exp_p = shared[simd_lane_id * 2 + 1];
+        sum_exp_q = simd_sum(sum_exp_q);
+        sum_exp_p = simd_sum(sum_exp_p);
+
+        lse_p = max_p + metal::fast::log(sum_exp_p);
+        lse_q = max_q + metal::fast::log(sum_exp_q);
+    }
+
+    threadgroup_barrier(mem_flags::mem_none);
+
+    {
+        float kl = 0;
+        float c = cotan[0];
+
+        int offset = thread_index_in_threadgroup * M;
+        for (int i = 0; i < full_blocks; i++) {
+            // Read and add to the kl
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = logits_q[offset + j];
+                vals_p[j] = logits_p[offset + j];
+            }
+
+            for (int j=0; j<M; j++) {
+                out[offset + j] = static_cast<T>(
+                    c * (metal::fast::exp(vals_q[j] - lse_q) - metal::fast::exp(vals_p[j] - lse_p)));
+            }
+
+            // Move to the next block
+            offset += block;
+        }
+        if (extra > 0) {
+            float vals_q[M];
+            float vals_p[M];
+            for (int j=0; j<M; j++) {
+                vals_q[j] = (offset + j < V) ? logits_q[offset + j] : -1e30;
+                vals_p[j] = (offset + j < V) ? logits_p[offset + j] : -1e30;
+            }
+
+            for (int j=0; j<M; j++) {
+                if (offset + j < V) {
+                    out[offset + j] = static_cast<T>(
+                        c * (metal::fast::exp(vals_q[j] - lse_q) - metal::fast::exp(vals_p[j] - lse_p)));
+                }
+            }
+        }
+    }
+    """
+
+    return mx.fast.metal_kernel(
+        name="kl_backward",
+        input_names=["logits_q", "logits_p", "cotan"],
+        output_names=["out"],
+        source=source,
+        ensure_row_contiguous=True,
+    )
+
+
+_kl_forward_kernel = _make_kl_forward_kernel()
+_kl_backward_kernel = _make_kl_backward_kernel()
+
+
+@mx.custom_function
+def _kl_div_loss(logits_q, logits_p):
+    n_outs = logits_q.size // logits_q.shape[-1]
+    dt = logits_q.dtype
+
+    return _kl_forward_kernel(
+        inputs=[logits_q, logits_p],
+        output_shapes=[logits_q.shape[:-1]],
+        output_dtypes=[dt],
+        template=[("T", dt), ("V", logits_q.shape[-1])],
+        grid=(1024, n_outs, 1),
+        threadgroup=(1024, 1, 1),
+    )[0]
+
+
+@_kl_div_loss.vjp
+def _kl_div_loss(primals, cotangent, output):
+    logits_q, logits_p = primals
+    dt = logits_q.dtype
+
+    dp = mx.zeros_like(logits_p)
+    dq = _kl_backward_kernel(
+        inputs=[logits_q, logits_p, cotangent],
+        output_shapes=[logits_q.shape],
+        output_dtypes=[dt],
+        template=[("T", dt), ("V", logits_q.shape[-1])],
+        grid=(1024, cotangent.size, 1),
+        threadgroup=(1024, 1, 1),
+    )[0]
+
+    return dq, dp
+
+
+def kl_div_loss(logits_q, logits_p):
+    if mx.metal.is_available():
+        return _kl_div_loss(logits_q, logits_p)
+    else:
+        return nn.losses.kl_div_loss(
+            logits_q - mx.logsumexp(logits_q, axis=-1, keepdims=True),
+            logits_p - mx.logsumexp(logits_p, axis=-1, keepdims=True),
+            axis=-1,
+            reduction="none",
+        )
@@ -1,20 +1,19 @@
 # Copyright © 2024 Apple Inc.

-import glob
-import shutil
+
 import time
 from dataclasses import dataclass, field
 from functools import partial
 from pathlib import Path
-from typing import List, Optional, Tuple

 import mlx.core as mx
 import mlx.nn as nn
 import numpy as np
 from mlx.nn.utils import average_gradients
 from mlx.utils import tree_flatten
-from transformers import PreTrainedTokenizer
+from tqdm import tqdm

+from .callbacks import TrainingCallback
 from .datasets import CacheDataset


@@ -93,7 +92,7 @@ def iterate_batches(
    if isinstance(dataset, CacheDataset):
        len_fn = lambda idx: dataset.itemlen(idx)
    else:
-        len_fn = lambda idx: len(dataset[idx])
+        len_fn = lambda idx: len(dataset[idx][0])
    idx = sorted(range(len(dataset)), key=len_fn)
    if len(dataset) < batch_size:
        raise ValueError(
@@ -164,13 +163,17 @@ def evaluate(

    index_iterator = iter(range(num_batches)) if num_batches != -1 else iter(int, 1)

-    for _, batch in zip(
-        index_iterator,
-        iterate_batches(
-            dataset=dataset,
-            batch_size=batch_size,
-            max_seq_length=max_seq_length,
+    for _, batch in tqdm(
+        zip(
+            index_iterator,
+            iterate_batches(
+                dataset=dataset,
+                batch_size=batch_size,
+                max_seq_length=max_seq_length,
+            ),
        ),
+        desc="Calculating loss...",
+        total=min(len(dataset) // batch_size, num_batches),
    ):
        losses, toks = loss(model, *batch)
        all_losses += losses * toks
@@ -183,17 +186,6 @@ def evaluate(
    return (all_losses / ntokens).item()


-class TrainingCallback:
-
-    def on_train_loss_report(self, train_info: dict):
-        """Called to report training loss at specified intervals."""
-        pass
-
-    def on_val_loss_report(self, val_info: dict):
-        """Called to report validation loss at specified intervals or the beginning."""
-        pass
-
-
 def train(
    model,
    optimizer,
@@ -274,7 +266,7 @@ def train(

            if training_callback is not None:
                val_info = {
-                    "iteration": it,
+                    "iteration": it - 1,
                    "val_loss": val_loss,
                    "val_time": val_time,
                }
@@ -54,6 +54,13 @@ def linear_to_lora_layers(
    """

    def to_lora(layer):
+        if not use_dora and hasattr(layer, "to_lora"):
+            return layer.to_lora(
+                r=config["rank"],
+                scale=config["scale"],
+                dropout=config["dropout"],
+            )
+
        if isinstance(layer, (nn.Linear, nn.QuantizedLinear)):
            LoRALayer = DoRALinear if use_dora else LoRALinear
        elif isinstance(layer, (SwitchLinear, QuantizedSwitchLinear)):
@@ -79,6 +86,7 @@ def linear_to_lora_layers(
        keys = set(keys)
    elif model.model_type in [
        "mistral",
+        "mistral3",
        "llama",
        "phi",
        "mixtral",
@@ -101,12 +109,15 @@ def linear_to_lora_layers(
        "cohere2",
        "minicpm",
        "minicpm3",
+        "minicpm4",
        "deepseek",
        "olmo2",
        "olmoe",
        "internlm3",
        "glm4",
        "mimo",
+        "ernie4_5",
+        "dots1",
    ]:
        keys = set(["self_attn.q_proj", "self_attn.v_proj"])
        if model.model_type in ["mixtral", "phimoe"]:
@@ -114,7 +125,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 in ["olmoe", "qwen3_moe"]:
+        if model.model_type in ["olmoe", "qwen3_moe", "dots1"]:
            keys.add("mlp.gate")

    elif model.model_type == "gpt_bigcode":
@@ -207,39 +218,36 @@ def dequantize(model: nn.Module) -> nn.Module:
    Returns:
        nn.Module: The model with dequantized layers.
    """
-    de_quantize_layers = []
+    dequantize_layers = []
    for name, module in model.named_modules():
+        bias = "bias" in module
        if isinstance(module, nn.QuantizedLinear):
-            bias = "bias" in module
-            weight = module.weight
-            weight = mx.dequantize(
-                weight,
-                module.scales,
-                module.biases,
-                module.group_size,
-                module.bits,
-            ).astype(mx.float16)
-            output_dims, input_dims = weight.shape
-            linear = nn.Linear(input_dims, output_dims, bias=bias)
-            linear.weight = weight
-            if bias:
-                linear.bias = module.bias
-            de_quantize_layers.append((name, linear))
-        if isinstance(module, nn.QuantizedEmbedding):
-            weight = mx.dequantize(
-                module.weight,
-                module.scales,
-                module.biases,
-                module.group_size,
-                module.bits,
-            ).astype(mx.float16)
-            num_embeddings, dims = weight.shape
-            emb = nn.Embedding(num_embeddings, dims)
-            emb.weight = weight
-            de_quantize_layers.append((name, emb))
+            cls = nn.Linear
+            kwargs = {"bias": bias}
+        elif isinstance(module, nn.QuantizedEmbedding):
+            kwargs = {}
+            cls = nn.Embedding
+        elif isinstance(module, QuantizedSwitchLinear):
+            kwargs = {"bias": bias}
+            cls = SwitchLinear
+        else:
+            continue
+        weight = mx.dequantize(
+            module.weight,
+            module.scales,
+            module.biases,
+            module.group_size,
+            module.bits,
+        )
+        args = weight.shape[::-1]
+        m = cls(*args, **kwargs)
+        if bias:
+            m.bias = module.bias
+        m.weight = weight
+        dequantize_layers.append((name, m))

-    if len(de_quantize_layers) > 0:
-        model.update_modules(tree_unflatten(de_quantize_layers))
+    if len(dequantize_layers) > 0:
+        model.update_modules(tree_unflatten(dequantize_layers))
    return model


@@ -262,20 +270,24 @@ def remove_lora_layers(model: nn.Module) -> nn.Module:
    return model


-def nparams(module):
-    if hasattr(module, "bits"):
-        n = 0 if not hasattr(module, "bias") else module.bias.size
-        return n + module.weight.size * 32 // module.bits
-    return sum(v.size for _, v in tree_flatten(module.parameters()))
-
-
-def print_trainable_parameters(model):
+def get_total_parameters(model):
    leaf_modules = tree_flatten(
        model.leaf_modules(), is_leaf=lambda m: isinstance(m, nn.Module)
    )
-    total_p = sum(nparams(m) for _, m in leaf_modules) / 10**6
+
+    def nparams(m):
+        if hasattr(m, "bits"):
+            n = 0 if not hasattr(m, "bias") else m.bias.size
+            return n + m.weight.size * 32 // m.bits
+        return sum(v.size for _, v in tree_flatten(m.parameters()))
+
+    return sum(nparams(m) for _, m in leaf_modules)
+
+
+def print_trainable_parameters(model):
+    total_p = get_total_parameters(model) / 1e6
    trainable_p = (
-        sum(v.size for _, v in tree_flatten(model.trainable_parameters())) / 10**6
+        sum(v.size for _, v in tree_flatten(model.trainable_parameters())) / 1e6
    )
    print(
        f"Trainable parameters: {(trainable_p * 100 / total_p):.3f}% "
@@ -3,6 +3,7 @@
 import copy
 import glob
 import importlib
+import inspect
 import json
 import logging
 import os
@@ -26,19 +27,17 @@ if os.getenv("MLXLM_USE_MODELSCOPE", "False").lower() == "true":
    try:
        from modelscope import snapshot_download
    except ImportError:
-        raise ImportError(
-            "Please run `pip install modelscope` to activate the ModelScope."
-        )
+        raise ImportError("Run `pip install modelscope` to use ModelScope.")
 else:
    from huggingface_hub import snapshot_download

-from mlx.utils import tree_flatten, tree_reduce
+from mlx.utils import tree_flatten, tree_map, tree_reduce
 from transformers import PreTrainedTokenizer

 # Local imports
 from .tokenizer_utils import TokenizerWrapper, load_tokenizer
 from .tuner.utils import dequantize as dequantize_model
-from .tuner.utils import load_adapters, nparams
+from .tuner.utils import get_total_parameters, load_adapters

 # Constants
 MODEL_REMAPPING = {
@@ -50,12 +49,6 @@ MODEL_REMAPPING = {
 MAX_FILE_SIZE_GB = 5


-class ModelNotFoundError(Exception):
-    def __init__(self, message):
-        self.message = message
-        super().__init__(self.message)
-
-
 def _get_classes(config: dict):
    """
    Retrieve the model and model args classes based on the configuration.
@@ -82,10 +75,7 @@ def compute_bits_per_weight(model):
    model_bytes = tree_reduce(
        lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, model, 0
    )
-    leaf_modules = tree_flatten(
-        model.leaf_modules(), is_leaf=lambda m: isinstance(m, nn.Module)
-    )
-    model_params = sum(nparams(m) for _, m in leaf_modules)
+    model_params = get_total_parameters(model)
    return model_bytes * 8 / model_params


@@ -99,37 +89,36 @@ def get_model_path(path_or_hf_repo: str, revision: Optional[str] = None) -> Path
        revision (str, optional): A revision id which can be a branch name, a tag, or a commit hash.

    Returns:
-        Path: The path to the model.
+        Tuple[Path, str]: A tuple containing the local file path and the Hugging Face repo ID.
    """
    model_path = Path(path_or_hf_repo)

    if not model_path.exists():
-        try:
-            model_path = Path(
-                snapshot_download(
-                    path_or_hf_repo,
-                    revision=revision,
-                    allow_patterns=[
-                        "*.json",
-                        "*.safetensors",
-                        "*.py",
-                        "tokenizer.model",
-                        "*.tiktoken",
-                        "tiktoken.model",
-                        "*.txt",
-                        "*.jsonl",
-                    ],
-                )
+        hf_path = path_or_hf_repo
+        model_path = Path(
+            snapshot_download(
+                path_or_hf_repo,
+                revision=revision,
+                allow_patterns=[
+                    "*.json",
+                    "*.safetensors",
+                    "*.py",
+                    "tokenizer.model",
+                    "*.tiktoken",
+                    "tiktoken.model",
+                    "*.txt",
+                    "*.jsonl",
+                    "*.jinja",
+                ],
            )
-        except:
-            raise ModelNotFoundError(
-                f"Model not found for path or HF repo: {path_or_hf_repo}.\n"
-                "Please make sure you specified the local path or Hugging Face"
-                " repo id correctly.\nIf you are trying to access a private or"
-                " gated Hugging Face repo, make sure you are authenticated:\n"
-                "https://huggingface.co/docs/huggingface_hub/en/guides/cli#huggingface-cli-login"
-            ) from None
-    return model_path
+        )
+    else:
+
+        from huggingface_hub import ModelCard
+
+        card = ModelCard.load(model_path / "README.md")
+        hf_path = card.data.base_model
+    return model_path, hf_path


 def load_config(model_path: Path) -> dict:
@@ -252,7 +241,7 @@ def load(
        FileNotFoundError: If config file or safetensors are not found.
        ValueError: If model class or args class are not found.
    """
-    model_path = get_model_path(path_or_hf_repo)
+    model_path, _ = get_model_path(path_or_hf_repo)

    model, config = load_model(model_path, lazy)
    if adapter_path is not None:
@@ -380,17 +369,18 @@ def upload_to_hub(path: str, upload_repo: str):
    print(f"Upload successful, go to https://huggingface.co/{upload_repo} for details.")


-def save_weights(
+def save_model(
    save_path: Union[str, Path],
-    weights: Dict[str, Any],
+    model: nn.Module,
    *,
-    donate_weights: bool = False,
+    donate_model: bool = False,
 ) -> None:
-    """Save model weights into specified directory."""
+    """Save model weights and metadata index into specified directory."""
    if isinstance(save_path, str):
        save_path = Path(save_path)
    save_path.mkdir(parents=True, exist_ok=True)

+    weights = dict(tree_flatten(model.parameters()))
    shards = make_shards(weights)
    shards_count = len(shards)
    shard_file_format = (
@@ -400,13 +390,20 @@ def save_weights(
    )

    total_size = sum(v.nbytes for v in weights.values())
-    index_data = {"metadata": {"total_size": total_size}, "weight_map": {}}
+    index_data = {
+        "metadata": {
+            "total_size": total_size,
+            "total_parameters": get_total_parameters(model),
+        },
+        "weight_map": {},
+    }
+    if donate_model:
+        model.update(tree_map(lambda _: mx.array([]), model.parameters()))

    # Write the weights and make sure no references are kept other than the
    # necessary ones
-    if donate_weights:
-        weights.clear()
-        del weights
+    weights.clear()
+    del weights

    for i in range(len(shards)):
        shard = shards[i]
@@ -456,8 +453,10 @@ def quantize_model(
            a dict of quantization parameters to pass to `to_quantized`.

    Returns:
-        Tuple: Tuple containing quantized weights and config.
+        Tuple: Tuple containing quantized model and config.
    """
+    if "quantization" in config:
+        raise ValueError("Cannot quantize already quantized model")
    quantized_config = copy.deepcopy(config)
    quantized_config["quantization"] = {"group_size": q_group_size, "bits": q_bits}

@@ -475,12 +474,11 @@ def quantize_model(
    )
    # support hf model tree #957
    quantized_config["quantization_config"] = quantized_config["quantization"]
-    quantized_weights = dict(tree_flatten(model.parameters()))

    bpw = compute_bits_per_weight(model)
    print(f"[INFO] Quantized model with {bpw:.3f} bits per weight.")

-    return quantized_weights, quantized_config
+    return model, quantized_config


 def save_config(
@@ -510,15 +508,15 @@ def save_config(
 def save(
    dst_path: Union[str, Path],
    src_path: Union[str, Path],
-    weights: Dict[str, mx.array],
+    model: nn.Module,
    tokenizer: TokenizerWrapper,
    config: Dict[str, Any],
    hf_repo: Optional[str] = None,
-    donate_weights: bool = True,
+    donate_model: bool = True,
 ):
    src_path = Path(src_path)
    dst_path = Path(dst_path)
-    save_weights(dst_path, weights, donate_weights=True)
+    save_model(dst_path, model, donate_model=True)
    save_config(config, config_path=dst_path / "config.json")
    tokenizer.save_pretrained(dst_path)

@@ -554,3 +552,18 @@ def common_prefix_len(list1, list2):
    # 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
+
+
+def does_model_support_input_embeddings(model: nn.Module) -> bool:
+    """
+    Check if the model supports input_embeddings in its call signature.
+    Args:
+        model (nn.Module): The model to check.
+    Returns:
+        bool: True if the model supports input_embeddings, False otherwise.
+    """
+    try:
+        signature = inspect.signature(model.__call__)
+        return "input_embeddings" in signature.parameters
+    except (ValueError, TypeError):
+        return False
@@ -24,17 +24,18 @@ setup(
    url="https://github.com/ml-explore/mlx-lm",
    license="MIT",
    install_requires=requirements,
-    packages=["mlx_lm", "mlx_lm.models", "mlx_lm.tuner"],
+    packages=["mlx_lm", "mlx_lm.models", "mlx_lm.quant", "mlx_lm.tuner"],
    python_requires=">=3.8",
    extras_require={
        "test": ["datasets"],
        "evaluate": ["lm-eval", "tqdm"],
-        "lwq": ["datasets"],
+        "quant": ["datasets", "tqdm"],
    },
    entry_points={
        "console_scripts": [
-            "mlx_lm.awq = mlx_lm.awq:main",
-            "mlx_lm.dwq = mlx_lm.dwq:main",
+            "mlx_lm.awq = mlx_lm.quant.awq:main",
+            "mlx_lm.dwq = mlx_lm.quant.dwq:main",
+            "mlx_lm.dynamic_quant = mlx_lm.quant.dynamic_quant: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,7 +43,6 @@ setup(
            "mlx_lm.fuse = mlx_lm.fuse:main",
            "mlx_lm.generate = mlx_lm.generate:main",
            "mlx_lm.lora = mlx_lm.lora:main",
-            "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",
@@ -90,6 +90,67 @@ class TestGenerate(unittest.TestCase):
        # from the target model, and last two should be drafts
        self.assertEqual(drafted, [True, True, False, True, True])

+    def test_stream_generate_input_embeddings(self):
+        sampler = make_sampler(temp=0.0)  # determinate sampler
+
+        # get prompt embeddings
+        messages = [{"role": "user", "content": "Say 'TEST' and nothing else"}]
+        prompt = self.tokenizer.apply_chat_template(
+            messages, add_generation_prompt=True
+        )
+        prompt_embeddings = self.model.model.embed_tokens(prompt)
+
+        response = ""
+        for generation_result in stream_generate(
+            model=self.model,
+            tokenizer=self.tokenizer,
+            prompt=[],  # no prompt tokens passed
+            max_tokens=5,
+            sampler=sampler,
+            input_embeddings=prompt_embeddings,
+        ):
+            response += generation_result.text
+
+        self.assertEqual("TEST", response)
+
+    def test_stream_generate_input_embeddings_prefill(self):
+        sampler = make_sampler(temp=0.0)  # determinate sampler
+
+        # get prompt embeddings
+        messages = [{"role": "user", "content": "Say 'TEST' and nothing else"}]
+        prompt = self.tokenizer.apply_chat_template(
+            messages, add_generation_prompt=True
+        )
+        prompt_embeddings = self.model.model.embed_tokens(prompt)
+
+        # setup prompt progress callback to track batched prefill
+        num_prompt_processing_callbacks = 0
+
+        def progress_callback(processed: int, total: int) -> None:
+            nonlocal num_prompt_processing_callbacks
+            num_prompt_processing_callbacks += 1
+
+        # generate
+        prefill_step_size = 5
+        response = ""
+        for generation_result in stream_generate(
+            model=self.model,
+            tokenizer=self.tokenizer,
+            prompt=[],  # no prompt tokens passed
+            max_tokens=5,
+            sampler=sampler,
+            input_embeddings=prompt_embeddings,
+            prefill_step_size=prefill_step_size,
+            prompt_progress_callback=progress_callback,
+        ):
+            response += generation_result.text
+
+        self.assertEqual("TEST", response)
+        num_embeddings = prompt_embeddings.shape[0]
+        self.assertEqual(
+            num_embeddings / prefill_step_size, num_prompt_processing_callbacks
+        )
+

 if __name__ == "__main__":
    unittest.main()
@@ -1,4 +1,5 @@
 # Copyright © 2024 Apple Inc.
+import copy
 import unittest

 import mlx.core as mx
@@ -230,6 +231,9 @@ class TestModels(unittest.TestCase):
            self.assertEqual(outputs.shape, (1, 1, vocab_size))
            self.assertEqual(outputs.dtype, t)

+        # Make sure the model can be copied / pickled
+        copy.deepcopy(model)
+
    def test_llama(self):
        from mlx_lm.models import llama

@@ -247,6 +251,24 @@ class TestModels(unittest.TestCase):
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

+    def test_bitnet(self):
+        from mlx_lm.models import bitnet
+
+        args = bitnet.ModelArgs(
+            model_type="bitnet",
+            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,
+        )
+        model = bitnet.Model(args)
+        self.model_test_runner(
+            model, args.model_type, args.vocab_size, args.num_hidden_layers
+        )
+
    def test_phi2(self):
        from mlx_lm.models import phi

@@ -256,15 +278,6 @@ class TestModels(unittest.TestCase):
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

-    def test_phixtral(self):
-        from mlx_lm.models import phixtral
-
-        args = phixtral.ModelArgs(
-            "phixtral", num_vocab=1000, num_layers=4, model_dim=1024
-        )
-        model = phixtral.Model(args)
-        self.model_test_runner(model, args.model_type, args.num_vocab, args.num_layers)
-
    def test_phi3(self):
        from mlx_lm.models import phi3

@@ -28,6 +28,13 @@ class DummyModelProvider:
                "chat_template": None,
                "use_default_chat_template": False,
                "trust_remote_code": False,
+                "num_draft_tokens": 3,
+                "temp": 0.0,
+                "top_p": 1.0,
+                "top_k": 0,
+                "min_p": 0.0,
+                "max_tokens": 512,
+                "chat_template_args": {},
            },
        )

@@ -123,6 +130,38 @@ class TestServer(unittest.TestCase):
        self.assertIn("id", response_body)
        self.assertIn("choices", response_body)

+    def test_handle_chat_completions_with_null_tool_content(self):
+        url = f"http://localhost:{self.port}/v1/chat/completions"
+        chat_post_data = {
+            "model": "chat_model",
+            "max_tokens": 10,
+            "temperature": 0.7,
+            "top_p": 0.85,
+            "repetition_penalty": 1.2,
+            "messages": [
+                {"role": "user", "content": "what is 2+3?"},
+                {
+                    "role": "assistant",
+                    "content": None,
+                    "tool_calls": [
+                        {
+                            "type": "function",
+                            "id": "123",
+                            "function": {
+                                "name": "add",
+                                "arguments": '{"a": 2, "b": 3}',
+                            },
+                        }
+                    ],
+                },
+                {"role": "tool", "content": "5", "tool_call_id": "123"},
+            ],
+        }
+        response = requests.post(url, json=chat_post_data)
+        response_body = response.text
+        self.assertIn("id", response_body)
+        self.assertIn("choices", response_body)
+
    def test_handle_models(self):
        url = f"http://localhost:{self.port}/v1/models"
        response = requests.get(url)
@@ -350,7 +389,9 @@ class TestGetPromptCache(unittest.TestCase):
        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):
+    @patch("mlx_lm.server.trim_prompt_cache")
+    @patch("mlx_lm.server.can_trim_prompt_cache", return_value=True)
+    def test_identical_request_full_hit(self, mock_can_trim, mock_trim_cache):
        """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)
@@ -361,10 +402,9 @@ class TestGetPromptCache(unittest.TestCase):
        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, [])
+        mock_trim_cache.assert_called_once_with("existing_cache_obj", 1)
+        self.assertEqual(processed_prompt, [3])
        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."""
@@ -24,6 +24,7 @@ class TestTokenizers(unittest.TestCase):
                    "tokenizer_config.json",
                    "special_tokens_map.json",
                    "tokenizer.model",
+                    "chat_template.jinja",
                ],
            )
        )
@@ -94,6 +95,24 @@ class TestTokenizers(unittest.TestCase):

        self.assertEqual(detokenizer.last_segment, tokenizer.eos_token)

+    def test_tool_calling(self):
+        tokenizer_repo = "mlx-community/Qwen3-4B-4bit"
+        tokenizer = self.download_tokenizer(tokenizer_repo)
+        self.assertTrue(tokenizer.has_tool_calling)
+        self.assertEqual(tokenizer.tool_call_start, "<tool_call>")
+        self.assertEqual(tokenizer.tool_call_end, "</tool_call>")
+
+        tokenizer_repo = "mlx-community/Llama-3.2-1B-Instruct-4bit"
+        tokenizer = self.download_tokenizer(tokenizer_repo)
+        self.assertTrue(tokenizer.has_tool_calling, False)
+
+    def test_thinking(self):
+        tokenizer_repo = "mlx-community/Qwen3-4B-4bit"
+        tokenizer = self.download_tokenizer(tokenizer_repo)
+        self.assertTrue(tokenizer.has_thinking)
+        self.assertEqual(tokenizer.think_start, "<think>")
+        self.assertEqual(tokenizer.think_end, "</think>")
+

 if __name__ == "__main__":
    unittest.main()
@@ -68,7 +68,8 @@ class TestUtils(unittest.TestCase):
            vocab_size=10_000,
        )
        model = llama.Model(args)
-        weights, config = utils.quantize_model(model, {}, 64, 4)
+        model, config = utils.quantize_model(model, {}, 64, 4)
+        weights = dict(tree_flatten(model.parameters()))
        self.assertTrue("model.layers.2.mlp.up_proj.scales" in weights)
        self.assertTrue("model.layers.2.mlp.up_proj.biases" in weights)
        self.assertEqual(config["quantization"]["group_size"], 64)
@@ -77,7 +78,7 @@ class TestUtils(unittest.TestCase):
    def test_convert(self):
        mlx_path = os.path.join(self.test_dir, "mlx_model")

-        convert(HF_MODEL_PATH, mlx_path=mlx_path, quantize=True)
+        convert(HF_MODEL_PATH, mlx_path=mlx_path, quantize=False)
        model, _ = utils.load(mlx_path)
        self.assertTrue(isinstance(model.layers[0].mlp.up_proj, nn.QuantizedLinear))
        self.assertTrue(isinstance(model.layers[-1].mlp.up_proj, nn.QuantizedLinear))
@@ -87,8 +88,8 @@ class TestUtils(unittest.TestCase):
        convert(HF_MODEL_PATH, mlx_path=mlx_path, dtype="bfloat16")
        model, _ = utils.load(mlx_path)

-        self.assertEqual(model.layers[0].mlp.up_proj.weight.dtype, mx.bfloat16)
-        self.assertEqual(model.layers[-1].mlp.up_proj.weight.dtype, mx.bfloat16)
+        self.assertEqual(model.layers[0].mlp.up_proj.scales.dtype, mx.bfloat16)
+        self.assertEqual(model.layers[-1].mlp.up_proj.scales.dtype, mx.bfloat16)


 if __name__ == "__main__":
@@ -32,7 +32,7 @@ class TestLoadModelCustomGetClasses(unittest.TestCase):
        def custom_get_classes(config):
            return CustomQwenModel, CustomQwenConfig

-        model_path = get_model_path(HF_MODEL_PATH)
+        model_path, _ = get_model_path(HF_MODEL_PATH)
        model, _ = load_model(model_path, get_model_classes=custom_get_classes)

        self.assertIsInstance(model, CustomQwenModel)
@@ -41,7 +41,7 @@ class TestLoadModelCustomGetClasses(unittest.TestCase):
        self.assertTrue(hasattr(model, "qwenWeights"))

    def test_load_model_with_default_get_classes(self):
-        model_path = get_model_path(HF_MODEL_PATH)
+        model_path, _ = get_model_path(HF_MODEL_PATH)
        model, _ = load_model(model_path)

        self.assertIsInstance(model, Qwen2Model)
Author	SHA1	Message	Date
Angelos Katharopoulos	6eb9059ce6	Replace recursion with an explicit stack	2025-07-06 16:31:49 -07:00
Angelos Katharopoulos	39a389c654	Add DP search across multiple quant options	2025-07-06 16:31:47 -07:00
Angelos Katharopoulos	29b74d0f95	Fix the formating	2025-07-06 16:11:38 -07:00
Angelos Katharopoulos	93b907f5d5	Move dwq and dynamic_quant to kl-loss	2025-07-06 15:54:58 -07:00
Angelos Katharopoulos	ed92899d1d	Add custom kl loss to tuner losses	2025-07-04 19:18:31 -07:00
Prince Canuma	5fa62eb5f5	Add bitnet1.58 with custom metal kernel (#219 ) * add bitnet * update activation to relu2 * working bitnet * remove artifacts * remove logging * add custom post quant * fix dtype and add compile * fixed weight unpack * add custom kernel to avoid memory overhead * compile relu2 * fix weight scale * remove unused * add tests and update tuner utils * update acknowledgements * add kernel caching * add act_quant and set float16 as default dtype * use mx.add and move scaling to kernel * remove act quant * move bitlinear layers to separate file * feat: add falcon-e and other bitnet support * refactor: address comments * add support for 1.58bit N-bit quants * 43.85% speedup in generation performance (M3 max) * refactor utils * remove masking (2% gen speed improvement) * add quantization config * test llama bitnet * refactor apply_hf_quant * default threadgroup: 64 -> 32 * add comment * fix prompt processing perf * remove modulo * compile kernel in the constructor * Improve the bitnet kernel * remove benchmark * refactor bitlinear swap * format * remove llama changes * revert utils * faster + cleanup * not trainable * fix tests --------- Co-authored-by: younesbelkada <younes.belkada@tii.ae> Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com> Co-authored-by: Awni Hannun <awni@apple.com>	2025-07-02 16:34:07 -07:00
Gökdeniz Gülmez	e8f8729854	Adding support for rednote-hilab/dots.llm1.inst (#211 ) * initial * adding the code * update ackn. * make trainable * clean ups * fix ModelArgs * fix args * fix * fix router forward * fix sanitize * nits * adding default vals * fix Convert indices to int32 * renaming like the torch version * fixes --------- Co-authored-by: Awni Hannun <awni@apple.com>	2025-07-02 13:03:26 -07:00
Awni Hannun	e8c2cfce6a	Allow converting local models (#265 ) * allow converting local models * fix dequantization	2025-07-02 12:35:18 -07:00
John Mai	5431546b1e	Add Ernie4.5 (#263 ) * feat: Add Ernie4.5 * nits * fix --------- Co-authored-by: Awni Hannun <awni@apple.com>	2025-07-02 07:42:38 -07:00
Awni Hannun	c2f7facb66	Gemma3n text only support (#258 ) * enable tool use in the server and add an example using openai client (#217) * simplify and speedup * fix rope * [Gemma3n] Only create required KV cache layers (#260) This commit makes the following changes 1. Refactors the calculation of which layers are shared up to the LanguageModel class. Each layer now receives the cache it is meant to operate on, whether unique or shared. 2. Add a make_cache method that constructs the KV cache layers according to the config. It only creates up to the first shared layer index. This fixes a bug in the pre-fill logic, which would throw an error when attempting to evaluate KV cache layers that had not yet (and would never be) updated with keys / values. * fix quantization + nits * compile --------- Co-authored-by: will-lms <will@lmstudio.ai>	2025-06-30 20:03:00 -07:00
Awni Hannun	36d0d04ecd	allow models to be pickled + test (#261 )	2025-06-29 11:27:40 -07:00
Awni Hannun	e6dfe18344	enable tool use in the server and add an example using openai client (#217 )	2025-06-26 06:02:28 -07:00
muhtasham	74a47b1434	fix tokenizer empty string (#253 )	2025-06-25 07:16:19 -07:00
will-lms	d0ef4bcf17	Pipe input_embeddings through mistral3 model_type (#254 )	2025-06-25 07:16:07 -07:00
Awni Hannun	7c13b0defc	fix cast predicate (#243 )	2025-06-17 15:07:31 -07:00
Awni Hannun	d9bd78a4db	support cuda back-end (#241 )	2025-06-17 14:00:43 -07:00
Angelos Katharopoulos	19287dc922	Implementation of AFM in MLX (#232 )	2025-06-12 12:24:45 -07:00
Magic Yang	4a3b2a978f	feat: add MiniCPM4 model structure code minicpm4.py and minicpm4 mode… (#212 ) * feat: add MiniCPM4 model structure code minicpm4.py and minicpm4 model_type * revert: remove .gitignore changes from PR * fix and cleanup su rope --------- Co-authored-by: Awni Hannun <awni@apple.com>	2025-06-09 07:58:21 -07:00
Awni Hannun	f009881e5c	fix dynamic quant for bias (#216 )	2025-06-08 21:18:16 -07:00
Awni Hannun	584780a05f	fix quant package (#214 )	2025-06-07 08:40:42 -07:00
Awni Hannun	e673a97c80	some nits in lora (#208 )	2025-06-05 16:33:16 -07:00
Awni Hannun	3be51537a3	Dynamic quants (#202 ) * dynamic quants + reorg * readme * angelos fix * Change sensitivity metric * update version * fix rebase --------- Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>	2025-06-02 12:31:42 -07:00
Awni Hannun	19153e1671	Add total params to metadata + cleanup (#207 ) * add total params to metadata + cleanup * comments	2025-06-02 12:08:06 -07:00
Awni Hannun	1db99d41a2	fix chat templates (#205 )	2025-06-01 08:50:04 -07:00
Gökdeniz Gülmez	d1d0771e3f	Update LoRA training logic to reset lora_parameters when fine-tuning with full model (#200 )	2025-05-29 07:29:07 -07:00
Gökdeniz Gülmez	e8980c050b	nits for qwen3_moe file (#199 ) * nits * remove unused initialize_rope, math inputs * remove Literal import in req_gemma * remove unused math import * remove unused Tuple import in phi * remove unused Tuple import in openlm * remove unused Tuple import in olmo * remove unused dataclass_fields import in nemotron * __"__ * __"__ * add apple compyrigth to mimo * more clean up * more nits	2025-05-28 21:46:59 -07:00
Awni Hannun	3cc61aa64d	DWQ updates (#196 ) * dwq nits * remove temp * fix sort * padding	2025-05-28 21:13:31 -07:00
Ryan	77edf17bc0	Add Evaluation Progress (#158 ) * Add Evaluation Progress Adding a tqdm progress bar for evaluate to make it more transparent what’s happening. * Add flag to suppress progress Add a quick boolean to suppress displaying the evaluation by default, but show it for —test. * nits / simplify * fix --------- Co-authored-by: Awni Hannun <awni@apple.com>	2025-05-23 06:44:36 -07:00
Awni Hannun	71e8e57c2e	Faster top-p and min-p sampling (#187 ) * Faster top-p and min-p sampling * comment	2025-05-19 16:55:38 -07:00
Ryan	1b555aaa08	Add the correct directory for LoRA Data (#188 )	2025-05-19 15:52:58 -07:00
Jesse Pangburn	77898fd22d	fix server.py null tool content (#185 ) (#186 ) * fix server.py null tool content (#185) * move null tool content fix into process_message_content (#185)	2025-05-19 11:20:01 -07:00
Gökdeniz Gülmez	f2aa9419d9	fix (#184 )	2025-05-19 07:26:43 -07:00
Awni Hannun	064c75d78e	fix server cache (#183 )	2025-05-17 19:22:09 -07:00
Matt Clayton	0824576a57	Pixtral text support, pipe input_embeddings through llama arch (#181 )	2025-05-16 19:15:00 -07:00
Ian	5960ee9c7a	fix upload to hub usage (#180 )	2025-05-16 06:45:01 -07:00
Gökdeniz Gülmez	29f8e7765d	Wandb fix (#177 ) * fix * fix * Enhance WandBCallback to log training and validation information with iteration steps. This improves the tracking of training progress in Weights & Biases.	2025-05-14 20:56:57 -07:00
Matt Clayton	f93589cb7d	Add input_embeddings input to generate_step, Gemma 3, Qwen 2 (#179 ) * Add input_embeddings input to generate_step, Gemma 3, Qwen 2 * PR responses * Better docstring for input_embeddings	2025-05-14 11:54:51 -07:00
Awni Hannun	864f5ce118	version bump	2025-05-14 06:19:52 -07:00
Awni Hannun	5101aebe05	Don't hide exception with wrong message when HF model isn't found (#173 ) * don't hide exception with wrong message * don't hide exception with wrong message * more nits * fix test	2025-05-13 13:57:33 -07:00
Ivan Fioravanti	1ca5474822	Update LORA.md for wikisql (#176 ) Fix for #174	2025-05-13 11:51:02 -07:00
Bence Szalai	4401043b0c	Do not pass empty tools array to the tokenizer (#170 ) * Do not pass empty tools array to the tokenizer * Update mlx_lm/server.py --------- Co-authored-by: Awni Hannun <awni.hannun@gmail.com>	2025-05-13 10:18:22 -07:00
rrziehe	76c30edbd4	add default params to better run quen3 (#162 ) * add default param to better run quen3 * code formatting * Refactor APIHandler to use chat template arguments and remove disable thinking option * fix test	2025-05-13 06:39:30 -07:00
Awni Hannun	854c580f72	DWQ extra activation penalty (#161 ) * Add L1 activation penalty for DWQ * Update defaults + layer selection * revert default temp * add some tips to the docs * nits	2025-05-12 17:21:52 -07:00
prldev	2973b75c8a	Allow to use mistral3 model type for lora (#172 )	2025-05-12 06:13:10 -07:00
Gökdeniz Gülmez	4b484773cf	adding report-to-wandb (#9 ) * update lora_config.yaml + LORA.md + lora.py * code formatting * udpaet Acknowledgements.md * nits * Refactor WandB integration in lora.py and trainer.py - Updated WandB reporting mechanism to use a project name argument instead of a boolean flag. - Removed the old TrainingCallback class definition from trainer.py and imported it from callbacks. - Adjusted argument parsing to accommodate the new WandB configuration. * Enhance WandBCallback to include log directory in initialization - Added log_dir parameter to WandBCallback constructor for specifying the logging directory. - Updated lora.py to pass adapter_path as log_dir when initializing WandBCallback. * nits * formating * README.md * update example yaml * nits * nits * nits in readme --------- Co-authored-by: Awni Hannun <awni@apple.com>	2025-05-09 19:43:11 -07:00
Awni Hannun	f1572d4586	mistral3 (#166 )	2025-05-09 19:42:53 -07:00
cblomert	c592f76f6a	Fixed calculation of n_dims for odd bit quants (#165 )	2025-05-09 06:17:43 -07:00
Yi Hong Ang	02a0241581	Add rudimentary health check (#155 )	2025-05-08 00:35:53 -07:00