Fix Nemotron H loading error (#426)

* fix

* format

.circleci/config.yml

@@ -20,7 +20,7 @@ jobs:
   mlx_lm_build_and_test:
     macos:
       xcode: "15.2.0"
-    resource_class: macos.m1.large.gen1
+    resource_class: m2pro.medium
     steps:
       - checkout
       - run:
@@ -30,6 +30,7 @@ jobs:
             python3.9 -m venv env
             source env/bin/activate
             pip install --upgrade pip
+            pip install sentencepiece
             pip install unittest-xml-reporting
             pip install -e ".[test]"
       - run:
@@ -40,6 +41,30 @@ jobs:
       - store_test_results:
           path: test-results
 
+  build_release:
+    macos:
+      xcode: "15.2.0"
+    resource_class: m2pro.medium
+    steps:
+      - checkout
+      - run:
+          name: Install dependencies
+          command: |
+            brew install python@3.9
+            python3.9 -m venv env
+            source env/bin/activate
+            pip install --upgrade pip
+            pip install build
+            pip install twine
+      - run:
+          name: Build and upload
+          command: |
+            source env/bin/activate
+            python -m build
+            twine upload dist/*
+      - store_artifacts:
+          path: dist/
+
 workflows:
   build_and_test:
     when:
@@ -50,6 +75,15 @@ workflows:
       - mlx_lm_build_and_test
       - linux_build_and_test
 
+  build_pypi_release:
+    jobs:
+      - build_release:
+          filters:
+            tags:
+              only: /^v.*/
+            branches:
+              ignore: /.*/
+
   prb:
     when:
       matches:

ACKNOWLEDGMENTS.md

@@ -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`, and Allenai's `OLMoE`; Added support for the following training algorithms: `full-fine-tuning`.
+- Gökdeniz Gülmez: Added support for the following architectures: OpenBMB's `MiniCPM` and `MiniCPM3`, Kyutai's `Helium`, State-Space's`Mamba v1`, Z.ai & THUKEG's `GLM4`, Rednote `dots.llm1`, Baisu's `Ernie4.5 MoE`, inclusionAI's `Bailing MoE e.g. Ling-family`, IBM's `Granite MoE`, Meituan's `LongCat`, Nvidia's `Nemotron H`, Swiss-AI's `Apertus`, and Allenai's `OLMoE`; Added support for the following training algorithms: `Full Weight Fine-Tuning`, and the `Muon` optimizer; 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`.

README.md

@@ -12,6 +12,7 @@ Some key features include:
   fine-tuning](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/LORA.md)
   with support for quantized models.
 * Distributed inference and fine-tuning with `mx.distributed`
+* Supports experiment tracking using SwanLab and W&B.
 
 The easiest way to get started is to install the `mlx-lm` package:
 

mlx_lm/LEARNED_QUANTS.md

@@ -0,0 +1,170 @@
+# Learned Quantization 
+
+To reduce the quality loss from quantization MLX LM has several options:
+
+- Distilled Weight Quantization (DWQ)
+- Activation-aware Weight Quantization (AWQ)[^1]
+- Dynamic quantization
+- GPT Quantization (GPTQ)[^2]
+
+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. GPTQ finds quantized weights which
+minimize the squared error of each layer's output given the provided input.
+
+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[train]"
+```
+
+### DWQ
+
+Use `mlx_lm.dwq` to run DWQ on a given model. For example:
+
+```bash
+mlx_lm.dwq --model Qwen/Qwen3-0.6B
+```
+
+Some important options, along with their default values are:
+
+- `--mlx-path mlx_model`: The location to save the DWQ model.
+- `--bits 4`: Precision of the quantization.
+- `--num-samples 1024`: Number of samples to use. Using more samples can lead to
+  better results but takes longer.
+- `--batch-size 8`: Use a smaller batch size to reduce the memory footprint.
+
+For a full list of options run:
+
+```bash
+mlx_lm.dwq --help
+```
+
+#### 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 Qwen/Qwen3-0.6B
+```
+
+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:
+
+```bash
+mlx_lm.awq --model Qwen/Qwen3-0.6B
+```
+
+The script can take anywhere form a few minutes to several hours to run
+depending on the model size and the number of samples.
+
+Some important options, along with their default values, are:
+
+- `--mlx-path mlx_model`: The location to save the AWQ model.
+- `--bits 4`: Precision of the quantization.
+- `--num-samples 32`: Number of samples to use. Using more samples can lead to
+  better results but takes longer.
+- `--n-grid 10`: The granularity of the AWQ search. A larger grid can lead to
+  better results but takes longer.
+
+For a full list of options run:
+
+```bash
+mlx_lm.awq --help
+```
+
+### GPTQ
+
+Use `mlx_lm.gptq` to run GPTQ on a given model. For example:
+
+```bash
+mlx_lm.awq --model Qwen/Qwen3-0.6B
+```
+
+The script can take anywhere from a few minutes to several hours depending on
+the model size.
+
+Some important options, along with their default values, are:
+
+- `--mlx-path mlx_model`: The location to save the AWQ model.
+- `--bits 4`: Precision of the quantization.
+
+
+### Evaluate
+
+Once the quantization training finishes, you can evaluate the quality of the
+model on downstream tasks using `mlx_lm.evaluate`. For example:
+
+```bash
+mlx_lm.evaluate \
+    --model mlx_model \
+    --tasks winogrande boolq arc_challenge arc_easy hellaswag openbookqa piqa social_iqa                     
+```
+
+### Upload to Hugging Face
+
+Use `mlx_lm.upload` to upload the quantized model to the Hugging Face Hub. For
+example:
+
+```bash
+mlx_lm.upload \
+    --path mlx_model \
+    --upload-repo mlx-community/Mistral-7B-Instruct-v0.3-3bit-DWQ
+```
+
+[^1]: Refer to the [paper](https://arxiv.org/abs/2306.00978)
+and [github repository](https://github.com/mit-han-lab/llm-awq) for more
+details on AWQ.
+[^2]: Refer to the [paper](https://arxiv.org/abs/2210.17323) for more details
+on GPTQ.

mlx_lm/LORA.md

@@ -26,6 +26,12 @@ LoRA (QLoRA).[^qlora] LoRA fine-tuning works with the following model families:
 
 ## Run
 
+First, make sure you have the training dependenices installed:
+
+```shell
+pip install "mlx-lm[train]"
+```
+
 The main command is `mlx_lm.lora`. To see a full list of command-line options run:
 
 ```shell
@@ -76,6 +82,17 @@ 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 using `--report-to wandb`, or
+to SwanLab using `--report-to swanlab`. Make sure to install the required
+packages beforehand: `pip install wandb` or `pip install swanlab`. You can
+enable both tracking tools simultaneously by separating them with a comma, for
+example: `--report-to wandb,swanlab`.
+
+To specify a project name for the logging tracker, use `--project-name <YOUR
+PROJECT NAME>`.
+
 #### Prompt Masking
 
 The default training computes a loss for every token in the sample. You can
@@ -291,7 +308,7 @@ example:
 
 ```yaml
 hf_dataset:
-  name: "billsum"
+  path: "billsum"
   prompt_feature: "text"
   completion_feature: "summary"
 ```
@@ -308,12 +325,12 @@ with the same structure as above. For example:
 
 ```yaml
 hf_dataset:
-  - name: "Open-Orca/OpenOrca"
+  - path: "Open-Orca/OpenOrca"
     train_split: "train[:90%]"
     valid_split: "train[-10%:]"
     prompt_feature: "question"
     completion_feature: "response"
-  - name: "trl-lib/ultrafeedback_binarized"
+  - path: "trl-lib/ultrafeedback_binarized"
     train_split: "train[:90%]"
     valid_split: "train[-10%:]"
     chat_feature: "chosen"
@@ -379,7 +396,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

mlx_lm/MERGE.md

@@ -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.

mlx_lm/SERVER.md

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

mlx_lm/UPLOAD.md

@@ -1,37 +0,0 @@
-### Packaging for PyPI
-
-Install `build` and `twine`:
-
-```
-pip install --user --upgrade build
-pip install --user --upgrade twine
-```
-
-Generate the source distribution and wheel:
-
-```
-python -m build
-```
-
-> [!warning]
-> Use a test server first
-
-#### Test Upload
-
-Upload to test server:
-
-```
-python -m twine upload --repository testpypi dist/*
-```
-
-Install from test server and check that it works:
-
-```
-python -m pip install --index-url https://test.pypi.org/simple/ --no-deps mlx-lm
-```
-
-#### Upload
-
-```
-python -m twine upload dist/*
-```

mlx_lm/__main__.py

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

mlx_lm/_version.py

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

mlx_lm/benchmark.py

@@ -0,0 +1,106 @@
+# Copyright © 2025 Apple Inc.
+
+import argparse
+
+import mlx.core as mx
+
+from mlx_lm import stream_generate
+from mlx_lm.generate import DEFAULT_MODEL
+from mlx_lm.tokenizer_utils import load_tokenizer
+from mlx_lm.utils import (
+    fetch_from_hub,
+    get_model_path,
+)
+
+
+def setup_arg_parser():
+    """Set up and return the argument parser."""
+    parser = argparse.ArgumentParser(description="LLM benchmarking script")
+    parser.add_argument(
+        "--model",
+        type=str,
+        help=(
+            "The path to the local model directory or Hugging Face repo. "
+            f"If no model is specified, then {DEFAULT_MODEL} is used."
+        ),
+        default=None,
+    )
+    parser.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        help="Enable trusting remote code for tokenizer",
+    )
+    parser.add_argument(
+        "--prompt-tokens",
+        "-p",
+        default=512,
+        help="Length of prompt",
+        type=int,
+    )
+    parser.add_argument(
+        "--generation-tokens",
+        "-g",
+        default=1024,
+        help="Length of completion",
+        type=int,
+    )
+    parser.add_argument(
+        "--num-trials",
+        "-n",
+        default=5,
+        help="Number of timing trials",
+        type=int,
+    )
+    return parser
+
+
+def main():
+    parser = setup_arg_parser()
+    args = parser.parse_args()
+    mx.random.seed(0)
+
+    model_path = args.model or DEFAULT_MODEL
+
+    model_path, _ = get_model_path(model_path, revision=None)
+    model, config, _ = fetch_from_hub(model_path, trust_remote_code=True)
+    tokenizer = load_tokenizer(
+        model_path,
+        eos_token_ids=[],  # Empty to avoid early stopping
+        tokenizer_config_extra={"trust_remote_code": True},
+    )
+
+    prompt_tokens = args.prompt_tokens
+    generation_tokens = args.generation_tokens
+    prompt = mx.random.randint(0, config["vocab_size"], (prompt_tokens,))
+
+    def _bench():
+        for response in stream_generate(
+            model, tokenizer, prompt, max_tokens=generation_tokens
+        ):
+            pass
+        return response
+
+    print("Running warmup..")
+    _bench()
+
+    report_keys = ["prompt_tps", "generation_tps", "peak_memory"]
+    print(f"Timing with {prompt_tokens=} and {generation_tokens=}.")
+    responses = []
+    for i in range(args.num_trials):
+        response = _bench()
+        responses.append(response)
+        results = [(k, getattr(response, k)) for k in report_keys]
+        results = [f"{k}={v:.3f}" for k, v in results]
+        print(f"Trial {i+1}:  " + ", ".join(results))
+
+    def avg(k):
+        vals = (getattr(response, k) for response in responses)
+        return sum(vals) / args.num_trials
+
+    results = [(k, avg(k)) for k in report_keys]
+    results = [f"{k}={v:.3f}" for k, v in results]
+    print(f"Averages: " + ", ".join(results))
+
+
+if __name__ == "__main__":
+    main()

mlx_lm/cache_prompt.py

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

mlx_lm/chat.py

@@ -1,7 +1,6 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import argparse
-import json
 
 import mlx.core as mx
 
@@ -12,6 +11,8 @@ from .utils import load
 
 DEFAULT_TEMP = 0.0
 DEFAULT_TOP_P = 1.0
+DEFAULT_XTC_PROBABILITY = 0.0
+DEFAULT_XTC_THRESHOLD = 0.0
 DEFAULT_SEED = None
 DEFAULT_MAX_TOKENS = 256
 DEFAULT_MODEL = "mlx-community/Llama-3.2-3B-Instruct-4bit"
@@ -26,6 +27,11 @@ def setup_arg_parser():
         help="The path to the local model directory or Hugging Face repo.",
         default=DEFAULT_MODEL,
     )
+    parser.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        help="Enable trusting remote code for tokenizer",
+    )
     parser.add_argument(
         "--adapter-path",
         type=str,
@@ -37,6 +43,18 @@ def setup_arg_parser():
     parser.add_argument(
         "--top-p", type=float, default=DEFAULT_TOP_P, help="Sampling top-p"
     )
+    parser.add_argument(
+        "--xtc-probability",
+        type=float,
+        default=DEFAULT_XTC_PROBABILITY,
+        help="Probability of XTC sampling to happen each next token",
+    )
+    parser.add_argument(
+        "--xtc-threshold",
+        type=float,
+        default=0.0,
+        help="Thresold the probs of each next token candidate to be sampled by XTC",
+    )
     parser.add_argument(
         "--seed",
         type=int,
@@ -56,6 +74,11 @@ def setup_arg_parser():
         default=DEFAULT_MAX_TOKENS,
         help="Maximum number of tokens to generate",
     )
+    parser.add_argument(
+        "--system-prompt",
+        default=None,
+        help="System prompt to be used for the chat template",
+    )
     return parser
 
 
@@ -69,7 +92,9 @@ def main():
     model, tokenizer = load(
         args.model,
         adapter_path=args.adapter_path,
-        tokenizer_config={"trust_remote_code": True},
+        tokenizer_config={
+            "trust_remote_code": True if args.trust_remote_code else None
+        },
     )
 
     def print_help():
@@ -91,14 +116,25 @@ def main():
         if query == "h":
             print_help()
             continue
-        messages = [{"role": "user", "content": query}]
+        messages = []
+        if args.system_prompt is not None:
+            messages.append({"role": "system", "content": args.system_prompt})
+        messages.append({"role": "user", "content": query})
         prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
         for response in stream_generate(
             model,
             tokenizer,
             prompt,
             max_tokens=args.max_tokens,
-            sampler=make_sampler(args.temp, args.top_p),
+            sampler=make_sampler(
+                args.temp,
+                args.top_p,
+                xtc_threshold=args.xtc_threshold,
+                xtc_probability=args.xtc_probability,
+                xtc_special_tokens=(
+                    tokenizer.encode("\n") + list(tokenizer.eos_token_ids)
+                ),
+            ),
             prompt_cache=prompt_cache,
         ):
             print(response.text, flush=True, end="")
@@ -106,4 +142,8 @@ def main():
 
 
 if __name__ == "__main__":
+    print(
+        "Calling `python -m mlx_lm.chat...` directly is deprecated."
+        " Use `mlx_lm.chat...` or `python -m mlx_lm chat ...` instead."
+    )
     main()

mlx_lm/convert.py

@@ -1,45 +1,64 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import argparse
-import glob
-import shutil
 from pathlib import Path
 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,
     fetch_from_hub,
     get_model_path,
     quantize_model,
-    save_config,
-    save_weights,
+    save,
     upload_to_hub,
 )
 
 
 def mixed_quant_predicate_builder(
-    low_bits: int = 4, high_bits: int = 4, group_size: int = 64
+    recipe: str, model: nn.Module
 ) -> Callable[[str, nn.Module, dict], Union[bool, dict]]:
+    high_bits = 6
+    group_size = 64
+
+    if recipe == "mixed_2_6":
+        low_bits = 2
+    elif recipe == "mixed_3_4":
+        low_bits = 3
+        high_bits = 4
+    elif recipe == "mixed_3_6":
+        low_bits = 3
+    elif recipe == "mixed_4_6":
+        low_bits = 4
+    else:
+        raise ValueError("Invalid quant recipe {recipe}")
+
+    down_keys = [k for k, _ in model.named_modules() if "down_proj" in k]
+    if len(down_keys) == 0:
+        raise ValueError("Model does not have expected keys for mixed quant.")
+
+    # Look for the layer index location in the path:
+    for layer_location, k in enumerate(down_keys[0].split(".")):
+        if k.isdigit():
+            break
+    num_layers = len(model.layers)
+
     def mixed_quant_predicate(
         path: str,
         module: nn.Module,
-        config: dict,
     ) -> Union[bool, dict]:
         """Implements mixed quantization predicates with similar choices to, for example, llama.cpp's Q4_K_M.
         Ref: https://github.com/ggerganov/llama.cpp/blob/917786f43d0f29b7c77a0c56767c0fa4df68b1c5/src/llama.cpp#L5265
         By Alex Barron: https://gist.github.com/barronalex/84addb8078be21969f1690c1454855f3
         """
-
-        if not hasattr(module, "to_quantized"):
-            return False
-
-        index = int(path.split(".")[2]) if len(path.split(".")) > 2 else 0
-
-        num_layers = config["num_hidden_layers"]
+        index = (
+            int(path.split(".")[layer_location])
+            if len(path.split(".")) > layer_location
+            else 0
+        )
         use_more_bits = (
             index < num_layers // 8
             or index >= 7 * num_layers // 8
@@ -57,19 +76,9 @@ def mixed_quant_predicate_builder(
     return mixed_quant_predicate
 
 
-QUANT_RECIPES = {
-    "mixed_2_6": mixed_quant_predicate_builder(low_bits=3, high_bits=6),
-    "mixed_3_6": mixed_quant_predicate_builder(low_bits=2, high_bits=6),
-}
+QUANT_RECIPES = ["mixed_2_6", "mixed_3_4", "mixed_3_6", "mixed_4_6"]
 
-
-def quant_args(arg):
-    if arg not in QUANT_RECIPES:
-        raise argparse.ArgumentTypeError(
-            f"Invalid q-recipe {arg!r}. Choose from: {list(QUANT_RECIPES.keys())}"
-        )
-    else:
-        return QUANT_RECIPES[arg]
+MODEL_CONVERSION_DTYPES = ["float16", "bfloat16", "float32"]
 
 
 def convert(
@@ -78,13 +87,15 @@ def convert(
     quantize: bool = False,
     q_group_size: int = 64,
     q_bits: int = 4,
-    dtype: str = "float16",
+    q_mode: str = "affine",
+    dtype: Optional[str] = None,
     upload_repo: str = None,
     revision: Optional[str] = None,
     dequantize: bool = False,
     quant_predicate: Optional[
-        Callable[[str, nn.Module, dict], Union[bool, dict]]
+        Union[Callable[[str, nn.Module, dict], Union[bool, dict]], str]
     ] = None,
+    trust_remote_code: bool = False,
 ):
     # Check the save path is empty
     if isinstance(mlx_path, str):
@@ -97,41 +108,60 @@ def convert(
         )
 
     print("[INFO] Loading")
-    model_path = get_model_path(hf_path, revision=revision)
-    model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
+    model_path, hf_path = get_model_path(hf_path, revision=revision)
+    model, config, tokenizer = fetch_from_hub(
+        model_path, lazy=True, trust_remote_code=trust_remote_code
+    )
 
-    weights = dict(tree_flatten(model.parameters()))
-    dtype = getattr(mx, dtype)
-    weights = {k: v.astype(dtype) for k, v in weights.items()}
+    if isinstance(quant_predicate, str):
+        quant_predicate = mixed_quant_predicate_builder(quant_predicate, model)
+
+    if dtype is None:
+        dtype = config.get("torch_dtype", None)
+    if dtype in MODEL_CONVERSION_DTYPES:
+        print("[INFO] Using dtype:", dtype)
+        dtype = getattr(mx, dtype)
+        cast_predicate = getattr(model, "cast_predicate", lambda _: True)
+
+        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, q_group_size, q_bits, quant_predicate=quant_predicate
+        model, config = quantize_model(
+            model,
+            config,
+            q_group_size,
+            q_bits,
+            mode=q_mode,
+            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_weights(mlx_path, weights, donate_weights=True)
-
-    py_files = glob.glob(str(model_path / "*.py"))
-    for file in py_files:
-        shutil.copy(file, mlx_path)
-
-    tokenizer.save_pretrained(mlx_path)
-
-    save_config(config, config_path=mlx_path / "config.json")
+    save(
+        mlx_path,
+        model_path,
+        model,
+        tokenizer,
+        config,
+        hf_repo=hf_path,
+    )
 
     if upload_repo is not None:
-        upload_to_hub(mlx_path, upload_repo, hf_path)
+        upload_to_hub(mlx_path, upload_repo)
 
 
 def configure_parser() -> argparse.ArgumentParser:
@@ -158,18 +188,26 @@ def configure_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--q-bits", help="Bits per weight for quantization.", type=int, default=4
     )
+    parser.add_argument(
+        "--q-mode",
+        help="The quantization mode.",
+        type=str,
+        default="affine",
+        choices=["affine", "mxfp4"],
+    )
     parser.add_argument(
         "--quant-predicate",
-        help=f"Mixed-bit quantization recipe. Choices: {list(QUANT_RECIPES.keys())}",
-        type=quant_args,
+        help=f"Mixed-bit quantization recipe.",
+        choices=QUANT_RECIPES,
+        type=str,
         required=False,
     )
     parser.add_argument(
         "--dtype",
-        help="Type to save the non-quantized parameters.",
+        help="Type to save the non-quantized parameters. Defaults to config.json's `torch_dtype` or the current model weights dtype.",
         type=str,
-        choices=["float16", "bfloat16", "float32"],
-        default="float16",
+        choices=MODEL_CONVERSION_DTYPES,
+        default=None,
     )
     parser.add_argument(
         "--upload-repo",
@@ -184,6 +222,12 @@ def configure_parser() -> argparse.ArgumentParser:
         action="store_true",
         default=False,
     )
+    parser.add_argument(
+        "--trust-remote-code",
+        help="Trust remote code when loading tokenizer.",
+        action="store_true",
+        default=False,
+    )
     return parser
 
 

mlx_lm/evaluate.py

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

mlx_lm/examples/lora_config.yaml

@@ -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,10 @@ val_batches: 25
 # Adam learning rate.
 learning_rate: 1e-5
 
+# Services to report logs to (comma-separated): wandb, swanlab, or both ('wandb,swanlab').
+# report_to: wandb,swanlab
+# project_name: "Your-awesome-mlx-project-name"
+
 # Number of training steps between loss reporting.
 steps_per_report: 10
 

mlx_lm/examples/openai_tool_use.py

@@ -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](**json.loads(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)

mlx_lm/examples/pipeline_generate.py

@@ -5,8 +5,7 @@ Run with:
 
 ```
 mlx.launch \
- --hostfile /path/to/hosts.txt \
- --backend mpi \
+ --hostfile /path/to/hosts.json \
  /path/to/pipeline_generate.py \
  --prompt "hello world"
 ```
@@ -19,6 +18,7 @@ https://ml-explore.github.io/mlx/build/html/usage/distributed.html).
 
 import argparse
 import json
+import resource
 from pathlib import Path
 
 import mlx.core as mx
@@ -28,6 +28,9 @@ from mlx.utils import tree_flatten
 from mlx_lm import load, stream_generate
 from mlx_lm.utils import load_model, load_tokenizer
 
+# Needed for 8 bit model
+resource.setrlimit(resource.RLIMIT_NOFILE, (2048, 4096))
+
 
 def download(repo: str, allow_patterns: list[str]) -> Path:
     return Path(
@@ -47,9 +50,9 @@ def shard_and_load(repo):
 
     # Lazy load and shard model to figure out
     # which weights we need
-    model, _ = load_model(model_path, lazy=True, strict=False)
+    model, config = load_model(model_path, lazy=True, strict=False)
 
-    group = mx.distributed.init(backend="mpi")
+    group = mx.distributed.init()
     rank = group.rank()
     model.model.pipeline(group)
 
@@ -65,7 +68,11 @@ def shard_and_load(repo):
     download(args.model, allow_patterns=local_files)
 
     # Load and shard the model, and load the weights
-    tokenizer = load_tokenizer(model_path)
+    tokenizer = load_tokenizer(
+        model_path,
+        {"trust_remote_code": True},
+        eos_token_ids=config.get("eos_token_id", None),
+    )
     model, _ = load_model(model_path, lazy=True, strict=False)
     model.model.pipeline(group)
     mx.eval(model.parameters())
@@ -98,7 +105,7 @@ if __name__ == "__main__":
     )
     args = parser.parse_args()
 
-    group = mx.distributed.init(backend="mpi")
+    group = mx.distributed.init()
     rank = group.rank()
 
     def rprint(*args, **kwargs):

mlx_lm/fuse.py

@@ -1,19 +1,14 @@
 import argparse
-import glob
-import shutil
 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,
     get_model_path,
-    save_config,
-    save_weights,
+    save,
     upload_to_hub,
 )
 
@@ -38,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.",
@@ -73,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:
@@ -89,23 +80,18 @@ def main() -> None:
     if args.de_quantize:
         print("De-quantizing model")
         model = dequantize(model)
-
-    weights = dict(tree_flatten(model.parameters()))
-
-    save_path = Path(args.save_path)
-
-    save_weights(save_path, weights)
-
-    py_files = glob.glob(str(model_path / "*.py"))
-    for file in py_files:
-        shutil.copy(file, save_path)
-
-    tokenizer.save_pretrained(save_path)
-
-    if args.de_quantize:
         config.pop("quantization", None)
 
-    save_config(config, config_path=save_path / "config.json")
+    save_path = Path(args.save_path)
+    save(
+        save_path,
+        model_path,
+        model,
+        tokenizer,
+        config,
+        hf_repo=hf_path,
+        donate_model=False,
+    )
 
     if args.export_gguf:
         model_type = config["model_type"]
@@ -113,18 +99,20 @@ 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:
-        hf_path = args.hf_path or (
-            args.model if not Path(args.model).exists() else None
-        )
         if hf_path is None:
             raise ValueError(
                 "Must provide original Hugging Face repo to upload local model."
             )
-        upload_to_hub(args.save_path, args.upload_repo, hf_path)
+        upload_to_hub(args.save_path, args.upload_repo)
 
 
 if __name__ == "__main__":
+    print(
+        "Calling `python -m mlx_lm.fuse...` directly is deprecated."
+        " Use `mlx_lm.fuse...` or `python -m mlx_lm fuse ...` instead."
+    )
     main()

mlx_lm/generate.py

@@ -26,18 +26,19 @@ 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
 DEFAULT_TEMP = 0.0
 DEFAULT_TOP_P = 1.0
 DEFAULT_MIN_P = 0.0
+DEFAULT_TOP_K = 0
+DEFAULT_XTC_PROBABILITY = 0.0
+DEFAULT_XTC_THRESHOLD = 0.0
 DEFAULT_MIN_TOKENS_TO_KEEP = 1
 DEFAULT_SEED = None
 DEFAULT_MODEL = "mlx-community/Llama-3.2-3B-Instruct-4bit"
@@ -60,6 +61,11 @@ def setup_arg_parser():
         ),
         default=None,
     )
+    parser.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        help="Enable trusting remote code for tokenizer",
+    )
     parser.add_argument(
         "--adapter-path",
         type=str,
@@ -104,6 +110,21 @@ def setup_arg_parser():
     parser.add_argument(
         "--min-p", type=float, default=DEFAULT_MIN_P, help="Sampling min-p"
     )
+    parser.add_argument(
+        "--top-k", type=int, default=DEFAULT_TOP_K, help="Sampling top-k"
+    )
+    parser.add_argument(
+        "--xtc-probability",
+        type=float,
+        default=DEFAULT_XTC_PROBABILITY,
+        help="Probability of XTC sampling to happen each next token",
+    )
+    parser.add_argument(
+        "--xtc-threshold",
+        type=float,
+        default=0.0,
+        help="Thresold the probs of each next token candidate to be sampled by XTC",
+    )
     parser.add_argument(
         "--min-tokens-to-keep",
         type=int,
@@ -198,29 +219,35 @@ 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.
     """
-    model_bytes = tree_reduce(
-        lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, model, 0
-    )
-    max_rec_size = mx.metal.device_info()["max_recommended_working_set_size"]
-    if model_bytes > 0.9 * max_rec_size:
-        model_mb = model_bytes // 2**20
-        max_rec_mb = max_rec_size // 2**20
-        print(
-            f"[WARNING] Generating with a model that requires {model_mb} MB "
-            f"which is close to the maximum recommended size of {max_rec_mb} "
-            "MB. This can be slow. See the documentation for possible work-arounds: "
-            "https://github.com/ml-explore/mlx-lm/tree/main#large-models"
+    if not mx.metal.is_available():
+        try:
+            yield
+        finally:
+            pass
+    else:
+        model_bytes = tree_reduce(
+            lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, model, 0
         )
-    old_limit = mx.set_wired_limit(max_rec_size)
-    try:
-        yield None
-    finally:
-        if streams is not None:
-            for s in streams:
-                mx.synchronize(s)
-        else:
-            mx.synchronize()
-        mx.set_wired_limit(old_limit)
+        max_rec_size = mx.metal.device_info()["max_recommended_working_set_size"]
+        if model_bytes > 0.9 * max_rec_size:
+            model_mb = model_bytes // 2**20
+            max_rec_mb = max_rec_size // 2**20
+            print(
+                f"[WARNING] Generating with a model that requires {model_mb} MB "
+                f"which is close to the maximum recommended size of {max_rec_mb} "
+                "MB. This can be slow. See the documentation for possible work-arounds: "
+                "https://github.com/ml-explore/mlx-lm/tree/main#large-models"
+            )
+        old_limit = mx.set_wired_limit(max_rec_size)
+        try:
+            yield
+        finally:
+            if streams is not None:
+                for s in streams:
+                    mx.synchronize(s)
+            else:
+                mx.synchronize()
+            mx.set_wired_limit(old_limit)
 
 
 @dataclass
@@ -280,6 +307,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.
@@ -304,14 +332,28 @@ 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 instead of or in
+          conjunction with prompt tokens. Default: ``None``.
 
     Yields:
         Tuple[mx.array, mx.array]: One token and a vector of log probabilities.
     """
+    if input_embeddings is not None:
+        if not does_model_support_input_embeddings(model):
+            raise ValueError("Model does not support input embeddings.")
+        elif len(prompt) > 0 and len(prompt) != len(input_embeddings):
+            raise ValueError(
+                f"When providing input_embeddings, their sequence length ({len(input_embeddings)}) "
+                f"must match the sequence length of the prompt ({len(prompt)}), or the "
+                "prompt must be empty."
+            )
+    elif len(prompt) == 0:
+        raise ValueError(
+            "Either input_embeddings or prompt (or both) must be provided."
+        )
 
-    y = prompt
     tokens = None
 
     # Create the KV cache for generation
@@ -320,8 +362,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)
 
@@ -334,37 +374,71 @@ def generate_step(
 
     sampler = sampler or (lambda x: mx.argmax(x, axis=-1))
 
-    def _step(y):
+    def _model_call(input_tokens: mx.array, input_embeddings: Optional[mx.array]):
+        if input_embeddings is not None:
+            return model(
+                input_tokens, cache=prompt_cache, input_embeddings=input_embeddings
+            )
+        else:
+            return model(input_tokens, cache=prompt_cache)
+
+    def _step(input_tokens: mx.array, input_embeddings: Optional[mx.array] = None):
+        nonlocal tokens
+
         with mx.stream(generation_stream):
-            logits = model(y[None], cache=prompt_cache)
+            logits = _model_call(
+                input_tokens=input_tokens[None],
+                input_embeddings=(
+                    input_embeddings[None] if input_embeddings is not None else None
+                ),
+            )
+
             logits = logits[:, -1, :]
 
-            if logits_processors:
-                nonlocal tokens
-                tokens = mx.concat([tokens, y]) if tokens is not None else y
-
+            if logits_processors and len(input_tokens) > 0:
+                tokens = (
+                    mx.concat([tokens, input_tokens])
+                    if tokens is not None
+                    else input_tokens
+                )
                 for processor in logits_processors:
                     logits = processor(tokens, logits)
 
             quantize_cache_fn(prompt_cache)
 
             logprobs = logits - mx.logsumexp(logits, keepdims=True)
-            y = sampler(logprobs)
-            return y, logprobs.squeeze(0)
+            sampled = sampler(logprobs)
+            return sampled, logprobs.squeeze(0)
 
     with mx.stream(generation_stream):
-        total_prompt_tokens = y.size
+        total_prompt_tokens = (
+            len(input_embeddings) if input_embeddings is not None else len(prompt)
+        )
         prompt_processed_tokens = 0
-        while y.size > prefill_step_size:
-            model(y[:prefill_step_size][None], cache=prompt_cache)
+        prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens)
+        while total_prompt_tokens - prompt_processed_tokens > 1:
+            n_to_process = min(prefill_step_size, prompt.size - 1)
+            _model_call(
+                input_tokens=prompt[:n_to_process][None],
+                input_embeddings=(
+                    input_embeddings[:n_to_process][None]
+                    if input_embeddings is not None
+                    else None
+                ),
+            )
             quantize_cache_fn(prompt_cache)
             mx.eval([c.state for c in prompt_cache])
+            prompt_processed_tokens += n_to_process
             prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens)
-            prompt_processed_tokens += prefill_step_size
-            y = y[prefill_step_size:]
+            prompt = prompt[n_to_process:]
+            input_embeddings = (
+                input_embeddings[n_to_process:]
+                if input_embeddings is not None
+                else input_embeddings
+            )
             mx.clear_cache()
 
-        y, logprobs = _step(y)
+        y, logprobs = _step(input_tokens=prompt, input_embeddings=input_embeddings)
 
     mx.async_eval(y, logprobs)
     n = 0
@@ -436,8 +510,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) :]
@@ -609,6 +681,7 @@ def stream_generate(
         )
     else:
         kwargs.pop("max_kv_size", None)
+        kwargs.pop("prompt_progress_callback", None)
         token_generator = speculative_generate_step(
             prompt, model, draft_model, **kwargs
         )
@@ -733,7 +806,7 @@ def main():
     tokenizer_config = (
         {} if not using_cache else json.loads(metadata["tokenizer_config"])
     )
-    tokenizer_config["trust_remote_code"] = True
+    tokenizer_config["trust_remote_code"] = True if args.trust_remote_code else None
 
     model_path = args.model
     if using_cache:
@@ -806,7 +879,16 @@ def main():
             raise ValueError("Draft model tokenizer does not match model tokenizer.")
     else:
         draft_model = None
-    sampler = make_sampler(args.temp, args.top_p, args.min_p, args.min_tokens_to_keep)
+    sampler = make_sampler(
+        args.temp,
+        args.top_p,
+        args.min_p,
+        args.min_tokens_to_keep,
+        top_k=args.top_k,
+        xtc_probability=args.xtc_probability,
+        xtc_threshold=args.xtc_threshold,
+        xtc_special_tokens=tokenizer.encode("\n") + list(tokenizer.eos_token_ids),
+    )
     response = generate(
         model,
         tokenizer,

mlx_lm/lora.py

@@ -1,10 +1,9 @@
-# Copyright © 2024 Apple Inc.
-
 import argparse
 import math
 import os
 import re
 import types
+import warnings
 from pathlib import Path
 
 import mlx.core as mx
@@ -13,8 +12,8 @@ import mlx.optimizers as optim
 import numpy as np
 import yaml
 
-from .tokenizer_utils import TokenizerWrapper
-from .tuner.datasets import load_dataset
+from .tuner.callbacks import get_reporting_callbacks
+from .tuner.datasets import CacheDataset, load_dataset
 from .tuner.trainer import TrainingArgs, TrainingCallback, evaluate, train
 from .tuner.utils import (
     build_schedule,
@@ -48,6 +47,9 @@ CONFIG_DEFAULTS = {
     "optimizer_config": {
         "adam": {},
         "adamw": {},
+        "muon": {},
+        "sgd": {},
+        "adafactor": {},
     },
     "data": "data/",
     "seed": 0,
@@ -67,8 +69,11 @@ CONFIG_DEFAULTS = {
     "config": None,
     "grad_checkpoint": False,
     "lr_schedule": None,
-    "lora_parameters": {"rank": 8, "alpha": 16, "dropout": 0.0, "scale": 10.0},
+    "lora_parameters": {"rank": 8, "dropout": 0.0, "scale": 20.0},
     "mask_prompt": False,
+    "wandb": None,  # will be deprecated in a future release
+    "report_to": None,
+    "project_name": None,
 }
 
 
@@ -104,9 +109,9 @@ def build_parser():
     parser.add_argument(
         "--optimizer",
         type=str,
-        choices=["adam", "adamw"],
+        choices=["adam", "adamw", "muon", "sgd", "adafactor"],
         default=None,
-        help="Optimizer to use for training: adam or adamw",
+        help="Optimizer to use for training: adam, adamw, sgd, or adafactor.",
     )
     parser.add_argument(
         "--mask-prompt",
@@ -180,6 +185,27 @@ def build_parser():
         help="Use gradient checkpointing to reduce memory use.",
         default=None,
     )
+    parser.add_argument(  # will be deprecated in a future release
+        "--wandb",
+        type=str,
+        default=None,
+        help=(
+            "The 'wandb' argument is deprecated and will be removed in a future release. "
+            "Use 'report_to: wandb' and 'project_name' in the configuration instead."
+        ),
+    )
+    parser.add_argument(
+        "--report-to",
+        type=str,
+        default=None,
+        help="Services to report logs to ('wandb', 'swanlab', or 'wandb,swanlab').",
+    )
+    parser.add_argument(
+        "--project-name",
+        type=str,
+        default=None,
+        help="Project name for logging. Defaults to the name of the root directory.",
+    )
     parser.add_argument("--seed", type=int, help="The PRNG seed")
     return parser
 
@@ -187,7 +213,6 @@ def build_parser():
 def train_model(
     args,
     model: nn.Module,
-    tokenizer: TokenizerWrapper,
     train_set,
     valid_set,
     training_callback: TrainingCallback = None,
@@ -203,6 +228,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(
@@ -245,11 +272,16 @@ def train_model(
 
     optimizer_name = args.optimizer.lower()
     optimizer_config = args.optimizer_config.get(optimizer_name, {})
-
     if optimizer_name == "adam":
         opt_class = optim.Adam
     elif optimizer_name == "adamw":
         opt_class = optim.AdamW
+    elif optimizer_name == "muon":
+        opt_class = optim.Muon
+    elif optimizer_name == "sgd":
+        opt_class = optim.SGD
+    elif optimizer_name == "adafactor":
+        opt_class = optim.Adafactor
     else:
         raise ValueError(f"Unsupported optimizer: {optimizer_name}")
 
@@ -258,20 +290,18 @@ def train_model(
     # Train model
     train(
         model=model,
-        tokenizer=tokenizer,
         args=training_args,
         optimizer=opt,
-        train_dataset=train_set,
-        val_dataset=valid_set,
+        train_dataset=CacheDataset(train_set),
+        val_dataset=CacheDataset(valid_set),
         training_callback=training_callback,
     )
 
 
-def evaluate_model(args, model: nn.Module, tokenizer: TokenizerWrapper, test_set):
+def evaluate_model(args, model: nn.Module, test_set):
     test_loss = evaluate(
         model=model,
-        dataset=test_set,
-        tokenizer=tokenizer,
+        dataset=CacheDataset(test_set),
         batch_size=args.batch_size,
         num_batches=args.test_batches,
         max_seq_length=args.max_seq_length,
@@ -284,9 +314,23 @@ def evaluate_model(args, model: nn.Module, tokenizer: TokenizerWrapper, test_set
 
 def run(args, training_callback: TrainingCallback = None):
     np.random.seed(args.seed)
+    if args.wandb is not None:
+        warnings.warn(
+            "The 'wandb' argument is deprecated and will be removed in a future release. "
+            "Use 'report_to: wandb' and 'project_name' in the configuration instead.",
+            DeprecationWarning,
+        )
+        args.report_to = "wandb"
+        args.project_name = args.wandb
+    training_callback = get_reporting_callbacks(
+        args.report_to,
+        project_name=args.project_name,
+        log_dir=args.adapter_path,
+        config=vars(args),
+    )
 
     print("Loading pretrained model")
-    model, tokenizer = load(args.model)
+    model, tokenizer = load(args.model, tokenizer_config={"trust_remote_code": True})
 
     print("Loading datasets")
     train_set, valid_set, test_set = load_dataset(args, tokenizer)
@@ -298,13 +342,13 @@ def run(args, training_callback: TrainingCallback = None):
 
     elif args.train:
         print("Training")
-        train_model(args, model, tokenizer, train_set, valid_set, training_callback)
+        train_model(args, model, train_set, valid_set, training_callback)
     else:
         raise ValueError("Must provide at least one of --train or --test")
 
     if args.test:
         print("Testing")
-        evaluate_model(args, model, tokenizer, test_set)
+        evaluate_model(args, model, test_set)
 
 
 def main():
@@ -330,4 +374,8 @@ def main():
 
 
 if __name__ == "__main__":
+    print(
+        "Calling `python -m mlx_lm.lora...` directly is deprecated."
+        " Use `mlx_lm.lora...` or `python -m mlx_lm lora ...` instead."
+    )
     main()

mlx_lm/manage.py

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

mlx_lm/merge.py

@@ -1,172 +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__":
-    main()

mlx_lm/models/afm7.py

@@ -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

mlx_lm/models/apertus.py

@@ -0,0 +1,222 @@
+# Copyright © 2023-2025 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
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    mlp_bias: bool
+    num_attention_heads: int
+    attention_bias: bool
+    rms_norm_eps: float
+    vocab_size: int
+    num_key_value_heads: int
+    max_position_embeddings: int
+    rope_theta: float
+    post_norm: bool
+    qk_norm: bool
+    tie_word_embeddings: bool
+    rope_traditional: bool = False
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+
+
+@partial(mx.compile, shapeless=True)
+def xielu(x, alpha_p, alpha_n, beta, eps):
+    alpha_p = nn.softplus(alpha_p)
+    alpha_n = beta + nn.softplus(alpha_n)
+    return mx.where(
+        x > 0,
+        alpha_p * mx.square(x) + beta * x,
+        (mx.expm1(mx.minimum(x, eps)) - x) * alpha_n + beta * x,
+    )
+
+
+class XieLU(nn.Module):
+    def __init__(
+        self,
+        alpha_p_init=0.8,
+        alpha_n_init=0.8,
+        beta=0.5,
+        eps=-1e-6,
+    ):
+        super().__init__()
+        alpha_p_tensor = mx.array(alpha_p_init)
+        alpha_n_tensor = mx.array(alpha_n_init - beta)
+        self.alpha_p = mx.log(mx.exp(alpha_p_tensor) - 1)
+        self.alpha_n = mx.log(mx.exp(alpha_n_tensor) - 1)
+
+        self.beta = mx.array(beta)
+        self.eps = mx.array(eps)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return xielu(x, self.alpha_p, self.alpha_n, self.beta, self.eps)
+
+
+class ApertusMLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.up_proj = nn.Linear(
+            args.hidden_size, args.intermediate_size, bias=args.mlp_bias
+        )
+        self.down_proj = nn.Linear(
+            args.intermediate_size, args.hidden_size, bias=args.mlp_bias
+        )
+        self.act_fn = XieLU()
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(self.act_fn(self.up_proj(x)))
+
+
+class ApertusAttention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.num_key_value_heads = args.num_key_value_heads
+
+        self.head_dim = args.hidden_size // args.num_attention_heads
+        self.scale = self.head_dim**-0.5
+
+        self.q_proj = nn.Linear(
+            args.hidden_size, args.num_attention_heads * self.head_dim, bias=False
+        )
+        self.k_proj = nn.Linear(
+            args.hidden_size, args.num_key_value_heads * self.head_dim, bias=False
+        )
+        self.v_proj = nn.Linear(
+            args.hidden_size, args.num_key_value_heads * self.head_dim, bias=False
+        )
+        self.o_proj = nn.Linear(
+            args.num_attention_heads * self.head_dim, args.hidden_size, bias=False
+        )
+
+        self.q_norm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
+        self.k_norm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
+
+        self.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            args.rope_traditional,
+            args.rope_scaling,
+            args.max_position_embeddings,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+        queries = self.q_norm(
+            queries.reshape(B, L, self.num_attention_heads, -1)
+        ).transpose(0, 2, 1, 3)
+        keys = self.k_norm(keys.reshape(B, L, self.num_key_value_heads, -1)).transpose(
+            0, 2, 1, 3
+        )
+        values = values.reshape(B, L, self.num_key_value_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 ApertusDecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.self_attn = ApertusAttention(args)
+        self.mlp = ApertusMLP(args)
+
+        self.attention_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.feedforward_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:
+        h = x + self.self_attn(self.attention_layernorm(x), mask, cache)
+        out = h + self.mlp(self.feedforward_layernorm(h))
+        return out
+
+
+class ApertusModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            ApertusDecoderLayer(args=args) for _ in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> 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=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 = ApertusModel(args)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        out = self.model(inputs, mask, cache)
+        return self.lm_head(out)
+
+    def sanitize(self, weights):
+        for k, v in weights.items():
+            if k.endswith("alpha_p") or k.endswith("alpha_n"):
+                weights[k] = v.squeeze()
+        return weights
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/baichuan_m1.py

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

mlx_lm/models/bailing_moe.py

@@ -0,0 +1,322 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    intermediate_size: int
+    max_position_embeddings: int
+    moe_intermediate_size: int
+    num_experts: int
+    num_shared_experts: int
+    norm_topk_prob: bool
+    num_attention_heads: int
+    num_experts_per_tok: int
+    num_hidden_layers: int
+    num_key_value_heads: int
+    rms_norm_eps: float
+    rope_theta: float
+    vocab_size: int
+    first_k_dense_replace: int
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    use_bias: bool = False
+    use_qkv_bias: bool = False
+    norm_head: bool = False
+    norm_softmax: bool = False
+
+
+class BailingMoeMLP(nn.Module):
+    def __init__(self, args: ModelArgs, intermediate_size: Optional[int] = None):
+        super().__init__()
+        self.intermediate_size = (
+            intermediate_size
+            if intermediate_size is not None
+            else args.intermediate_size
+        )
+
+        self.gate_proj = nn.Linear(
+            args.hidden_size, self.intermediate_size, bias=args.use_bias
+        )
+        self.down_proj = nn.Linear(
+            self.intermediate_size, args.hidden_size, bias=args.use_bias
+        )
+        self.up_proj = nn.Linear(
+            args.hidden_size, self.intermediate_size, bias=args.use_bias
+        )
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class BailingMoeAttention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.num_key_value_heads = args.num_key_value_heads
+        self.head_dim = args.hidden_size // self.num_attention_heads
+        self.scale = self.head_dim**-0.5
+
+        self.query_key_value = nn.Linear(
+            args.hidden_size,
+            (self.num_attention_heads + 2 * self.num_key_value_heads) * self.head_dim,
+            bias=args.use_qkv_bias,
+        )
+        self.dense = nn.Linear(
+            self.num_attention_heads * self.head_dim,
+            args.hidden_size,
+            bias=args.use_bias,
+        )
+
+        self.rope = initialize_rope(
+            self.head_dim,
+            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
+
+        qkv = self.query_key_value(x)
+
+        q_size = self.num_attention_heads * self.head_dim
+        kv_size = self.num_key_value_heads * self.head_dim
+        q, k, v = mx.split(qkv, [q_size, q_size + kv_size], axis=-1)
+
+        queries = q.reshape(B, L, self.num_attention_heads, self.head_dim).transpose(
+            0, 2, 1, 3
+        )
+        keys = k.reshape(B, L, self.num_key_value_heads, self.head_dim).transpose(
+            0, 2, 1, 3
+        )
+        values = v.reshape(B, L, self.num_key_value_heads, self.head_dim).transpose(
+            0, 2, 1, 3
+        )
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.dense(output)
+
+
+class BailingMoeGate(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.top_k = config.num_experts_per_tok
+        self.num_experts = config.num_experts
+        self.norm_topk_prob = config.norm_topk_prob
+        self.gating_dim = config.hidden_size
+
+        self.gate_proj = nn.Linear(self.gating_dim, self.num_experts, bias=False)
+
+    def __call__(self, hidden_states):
+        B, L, D = hidden_states.shape
+        x = hidden_states.reshape(-1, D)
+
+        logits = self.gate_proj(x)
+        scores = mx.softmax(logits, axis=-1, precise=True)
+
+        topk_idx = mx.argpartition(scores, kth=-self.top_k, axis=-1)[..., -self.top_k :]
+        topk_scores = mx.take_along_axis(scores, topk_idx, axis=-1)
+
+        if self.top_k > 1 and self.norm_topk_prob:
+            denom = mx.sum(topk_scores, axis=-1, keepdims=True)
+            topk_scores = topk_scores / mx.maximum(denom, 1e-9)
+
+        return topk_idx, topk_scores
+
+
+class BailingMoeSparseMoeBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.num_experts_per_tok = args.num_experts_per_tok
+
+        self.switch_mlp = SwitchGLU(
+            args.hidden_size,
+            args.moe_intermediate_size,
+            args.num_experts,
+            bias=args.use_bias,
+        )
+
+        self.gate = BailingMoeGate(config=args)
+
+        if args.num_shared_experts > 0:
+            self.shared_experts = BailingMoeMLP(
+                args=args,
+                intermediate_size=args.moe_intermediate_size * args.num_shared_experts,
+            )
+        else:
+            self.shared_experts = None
+
+    def __call__(self, hidden_states):
+        batch_size, seq_len, hidden_dim = hidden_states.shape
+
+        if self.shared_experts is not None:
+            identity = hidden_states
+
+        x = hidden_states.reshape(-1, hidden_dim)
+
+        expert_indices, expert_weights = self.gate(hidden_states)
+        expert_outputs = self.switch_mlp(x, expert_indices)
+
+        weighted_output = mx.sum(expert_outputs * expert_weights[..., None], axis=-2)
+        output = weighted_output.reshape(batch_size, seq_len, hidden_dim)
+
+        if self.shared_experts is not None:
+            output = output + self.shared_experts(hidden_states)
+
+        return output
+
+
+class BailingMoeDecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.attention = BailingMoeAttention(args)
+
+        self.mlp = (
+            BailingMoeSparseMoeBlock(args)
+            if (
+                args.num_experts is not None and layer_idx >= args.first_k_dense_replace
+            )
+            else BailingMoeMLP(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.attention(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        return h + r
+
+
+class BailingMoeModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.word_embeddings = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            BailingMoeDecoderLayer(args, layer_idx=i)
+            for i in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ):
+        h = self.word_embeddings(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.norm_head = args.norm_head
+        self.model_type = args.model_type
+        self.model = BailingMoeModel(args)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        h = self.model(inputs, mask, cache)
+        return self.lm_head(h)
+
+    def sanitize(self, weights):
+        if self.norm_head:
+            w = weights["lm_head.weight"]
+            dtype = w.dtype
+            weight_norm = (
+                mx.linalg.norm(w.astype(mx.float32), axis=0, keepdims=True) + 1e-7
+            )
+            weights["lm_head.weight"] = (w / weight_norm).astype(dtype)
+
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}"
+
+            if l >= self.args.first_k_dense_replace:
+                for m in ["gate_proj", "down_proj", "up_proj"]:
+                    for k in ["weight", "scales", "biases"]:
+                        if f"{prefix}.mlp.experts.0.{m}.{k}" in weights:
+                            to_join = [
+                                weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}")
+                                for e in range(self.args.num_experts)
+                            ]
+                            weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(
+                                to_join
+                            )
+
+                if f"{prefix}.mlp.gate.weight" in weights:
+                    gate_weight = weights.pop(f"{prefix}.mlp.gate.weight")
+                    weights[f"{prefix}.mlp.gate.gate_proj.weight"] = gate_weight
+
+                if f"{prefix}.mlp.gate.bias" in weights:
+                    gate_bias = weights.pop(f"{prefix}.mlp.gate.bias")
+                    weights[f"{prefix}.mlp.gate.gate_proj.bias"] = gate_bias
+
+        return weights
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if path.endswith("mlp.gate"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/base.py

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

mlx_lm/models/bitlinear_layers.py

@@ -0,0 +1,158 @@
+# Copyright © 2025 Apple Inc.
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.nn.layers.quantized import QuantizedLinear
+from mlx.utils import tree_flatten, tree_unflatten
+
+
+def bitnet_quantize(model, quantization_config: dict):
+    quantize_layers = []
+    modules_to_not_convert = quantization_config.get("modules_to_not_convert", [])
+    invert_weight_scales = (
+        quantization_config.get("linear_class", "") != "autobitlinear"
+    )
+
+    for name, module in tree_flatten(model.leaf_modules(), is_leaf=nn.Module.is_module):
+
+        # Replace nn.Linear layers, but skip any layer from the `modules_to_not_convert` list
+        if name not in modules_to_not_convert and isinstance(module, nn.Linear):
+            old_weight = module.weight
+            out_features, in_features = old_weight.shape
+            bias = "bias" in module
+            new_layer = BitLinear(
+                in_features,
+                out_features,
+                bias=bias,
+                invert_weight_scales=invert_weight_scales,
+            )
+            quantize_layers.append((name, new_layer))
+    if len(quantize_layers) > 0:
+        model.update_modules(tree_unflatten(quantize_layers))
+    return model
+
+
+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

mlx_lm/models/bitnet.py

@@ -0,0 +1,216 @@
+# 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

mlx_lm/models/cache.py

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

mlx_lm/models/cohere.py

@@ -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

mlx_lm/models/dbrx.py

@@ -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
 

mlx_lm/models/deepseek.py

@@ -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):

mlx_lm/models/deepseek_v2.py

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

mlx_lm/models/deepseek_v3.py

@@ -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
@@ -97,9 +97,7 @@ class DeepseekV3YarnRotaryEmbedding(nn.Module):
             scaling_factor, mscale_all_dim
         )
         freq_extra = base ** (mx.arange(0, dim, 2, dtype=mx.float32) / dim)
-        freq_inter = scaling_factor * base ** (
-            mx.arange(0, dim, 2, dtype=mx.float32) / dim
-        )
+        freq_inter = scaling_factor * freq_extra
         low, high = yarn_find_correction_range(
             beta_fast,
             beta_slow,
@@ -126,12 +124,6 @@ class DeepseekV3YarnRotaryEmbedding(nn.Module):
         )
 
 
-# A clipped silu to prevent fp16 from overflowing
-@partial(mx.compile, shapeless=True)
-def clipped_silu(x):
-    return mx.clip(x * mx.sigmoid(x), a_min=-100, a_max=100)
-
-
 class DeepseekV3Attention(nn.Module):
     def __init__(self, config: ModelArgs):
         super().__init__()
@@ -157,7 +149,7 @@ class DeepseekV3Attention(nn.Module):
             self.q_a_proj = nn.Linear(
                 self.hidden_size, self.q_lora_rank, bias=config.attention_bias
             )
-            self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank)
+            self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank, eps=1e-6)
             self.q_b_proj = nn.Linear(
                 self.q_lora_rank, self.num_heads * self.q_head_dim, bias=False
             )
@@ -167,7 +159,7 @@ class DeepseekV3Attention(nn.Module):
             self.kv_lora_rank + self.qk_rope_head_dim,
             bias=config.attention_bias,
         )
-        self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank)
+        self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank, eps=1e-6)
         self.kv_b_proj = nn.Linear(
             self.kv_lora_rank,
             self.num_heads
@@ -291,19 +283,22 @@ def group_expert_select(
 
     k = top_k
     scores = mx.sigmoid(gates.astype(mx.float32))
+    orig_scores = scores
     scores = scores + e_score_correction_bias
     scores = mx.unflatten(scores, axis=-1, shape=(n_group, -1))
     group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True)
     k = n_group - topk_group
     group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :]
-    scores = mx.put_along_axis(scores, group_idx, mx.array(0.0), axis=-2)
+    scores = mx.put_along_axis(
+        scores, mx.stop_gradient(group_idx), mx.array(0.0), axis=-2
+    )
     scores = mx.flatten(scores, -2, -1)
 
     k = top_k
     inds = mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k]
-    scores = mx.take_along_axis(scores, inds, axis=-1)
+    scores = mx.take_along_axis(orig_scores, inds, axis=-1)
     if top_k > 1 and norm_topk_prob:
-        denominator = scores.sum(axis=-1, keepdims=True) + 1e-20
+        denominator = scores.sum(axis=-1, keepdims=True)
         scores = scores / denominator
     scores = scores * routed_scaling_factor
 
@@ -345,7 +340,6 @@ class DeepseekV3MoE(nn.Module):
             config.hidden_size,
             config.moe_intermediate_size,
             config.n_routed_experts,
-            activation=clipped_silu,
         )
 
         self.gate = MoEGate(config)
@@ -437,8 +431,6 @@ class DeepseekV3Model(nn.Module):
 
         pipeline_rank = self.pipeline_rank
         pipeline_size = self.pipeline_size
-        # Hack to avoid time-outs during prompt-processing
-        dist_stream = mx.cpu if h.shape[1] > 1 else mx.gpu
         if mask is None:
             mask = create_attention_mask(h, cache)
 
@@ -448,19 +440,17 @@ class DeepseekV3Model(nn.Module):
         # Receive from the previous process in the pipeline
 
         if pipeline_rank < pipeline_size - 1:
-            h = mx.distributed.recv_like(h, (pipeline_rank + 1), stream=dist_stream)
+            h = mx.distributed.recv_like(h, (pipeline_rank + 1))
 
         for i in range(self.num_layers):
             h = self.layers[self.start_idx + i](h, mask, cache[i])
 
         # Send to the next process in the pipeline
         if pipeline_rank != 0:
-            h = mx.distributed.send(
-                h, (pipeline_rank - 1) % pipeline_size, stream=dist_stream
-            )
+            h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
 
         # Broadcast h while keeping it in the graph
-        h = mx.distributed.all_gather(h, stream=dist_stream)[: h.shape[0]]
+        h = mx.distributed.all_gather(h)[: h.shape[0]]
 
         return self.norm(h)
 
@@ -484,6 +474,7 @@ class Model(nn.Module):
 
     def sanitize(self, weights):
         def dequant(weight, scale_inv):
+            dtype = weight.dtype
             bs = 128  # block size
             m, n = weight.shape
             pad_bottom = (-m) % bs
@@ -492,11 +483,10 @@ class Model(nn.Module):
             weight = weight.reshape(
                 ((m + pad_bottom) // bs, bs, (n + pad_side) // bs, bs)
             )
-            scale_inv = scale_inv.astype(weight.dtype)
             weight = (weight * scale_inv[:, None, :, None]).reshape(
                 m + pad_bottom, n + pad_side
             )
-            return weight[:m, :n]
+            return weight[:m, :n].astype(dtype)
 
         # Dequantize
         new_weights = {}
@@ -533,3 +523,10 @@ class Model(nn.Module):
     @property
     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
+
+        return predicate

mlx_lm/models/dots1.py

@@ -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

mlx_lm/models/ernie4_5.py

@@ -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

mlx_lm/models/ernie4_5_moe.py

@@ -0,0 +1,291 @@
+# Copyright © 2023-2024 Apple Inc.
+
+from dataclasses import dataclass, field
+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
+from .switch_layers import SwitchGLU
+
+
+@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
+    num_hidden_layers: int
+    rms_norm_eps: float
+    vocab_size: int
+    rope_theta: float
+    use_bias: bool
+    tie_word_embeddings: bool
+    moe_num_experts: int
+    moe_layer_start_index: int = 0
+    moe_intermediate_size: int = 0
+    moe_capacity: list[int] = field(default_factory=list)
+    moe_k: int = 1
+    moe_layer_interval: int = 1
+    moe_use_aux_free: bool = False
+    moe_num_shared_experts: int = 0
+    moe_layer_end_index: Optional[int] = None
+    head_dim: Optional[int] = None
+    moe_gate_act: str = "softmax"
+
+
+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 Ernie4_5_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 Ernie4_5_MoeMLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.k = args.moe_k
+        self.moe_intermediate_size = (
+            args.moe_intermediate_size
+            if args.moe_intermediate_size
+            else args.intermediate_size
+        )
+
+        self.gate = nn.Linear(args.hidden_size, args.moe_num_experts, bias=False)
+
+        self.switch_mlp = SwitchGLU(
+            args.hidden_size,
+            self.moe_intermediate_size,
+            args.moe_num_experts,
+            bias=args.use_bias,
+        )
+
+        if getattr(args, "moe_num_shared_experts", 0) > 0:
+            shared_intermediate_size = (
+                args.moe_intermediate_size * args.moe_num_shared_experts
+                if getattr(args, "moe_intermediate_size", None)
+                else args.intermediate_size * args.moe_num_shared_experts
+            )
+            self.shared_experts = Ernie4_5_MLP(
+                args.hidden_size, shared_intermediate_size, args.use_bias
+            )
+        else:
+            self.shared_experts = None
+
+        if args.moe_gate_act == "softmax":
+            self.gate_act = nn.Softmax()
+        elif args.moe_gate_act == "sigmoid":
+            self.gate_act = nn.Sigmoid()
+        else:
+            raise ValueError(f"{args.moe_gate_act} is not supported.")
+
+    def __call__(self, x: mx.array) -> mx.array:
+        gates = self.gate(x)
+        gates = self.gate_act(gates.astype(mx.float32))
+
+        k = self.k
+        inds = mx.stop_gradient(mx.argpartition(-gates, kth=k - 1, axis=-1)[..., :k])
+        scores = mx.take_along_axis(gates, inds, axis=-1)
+
+        scores = scores / mx.maximum(scores.sum(axis=-1, keepdims=True), 1e-12)
+
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
+
+        if self.shared_experts is not None:
+            y = y + self.shared_experts(x)
+
+        return y
+
+
+class Ernie4_5_DecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.self_attn = Attention(args)
+
+        moe_layer_start_index = (
+            min(args.moe_layer_start_index)
+            if isinstance(args.moe_layer_start_index, (tuple, list))
+            else args.moe_layer_start_index
+        )
+
+        if args.moe_layer_end_index is None:
+            moe_layer_end_index = args.num_hidden_layers - 1
+        else:
+            moe_layer_end_index = (
+                max(args.moe_layer_end_index)
+                if isinstance(args.moe_layer_end_index, (tuple, list))
+                else args.moe_layer_end_index
+            )
+
+        if (
+            ((layer_idx + 1) % args.moe_layer_interval == 0)
+            and layer_idx >= moe_layer_start_index
+            and layer_idx <= moe_layer_end_index
+        ):
+            self.mlp = Ernie4_5_MoeMLP(args)
+        else:
+            self.mlp = Ernie4_5_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 = [
+            Ernie4_5_DecoderLayer(args, i) for i in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        h = self.embed_tokens(inputs)
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = 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
+
+    def sanitize(self, weights):
+        remove_patterns = [
+            "mtp_block.",
+            "mtp_linear_proj.",
+            "mtp_hidden_norm.",
+            "mtp_emb_norm.",
+            "e_score_correction_bias",
+        ]
+
+        weights = {
+            key: value
+            for key, value in weights.items()
+            if not any(pattern in key for pattern in remove_patterns)
+        }
+
+        # Stack experts
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}"
+            for m in ["gate_proj", "down_proj", "up_proj"]:
+                if f"{prefix}.mlp.experts.0.{m}.weight" in weights:
+                    to_join = [
+                        weights.pop(f"{prefix}.mlp.experts.{e}.{m}.weight")
+                        for e in range(self.args.moe_num_experts)
+                    ]
+                    weights[f"{prefix}.mlp.switch_mlp.{m}.weight"] = mx.stack(to_join)
+
+        return weights

mlx_lm/models/exaone4.py

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

mlx_lm/models/gemma.py

@@ -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

mlx_lm/models/gemma2.py

@@ -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

mlx_lm/models/gemma3.py

@@ -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()))

mlx_lm/models/gemma3_text.py

@@ -1,12 +1,13 @@
 # Copyright © 2025 Apple Inc.
 
 from dataclasses import dataclass
+from functools import partial
 from typing import Any, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .cache import KVCache, RotatingKVCache
 
 
@@ -88,9 +89,8 @@ class Attention(nn.Module):
         # Sliding window
         if isinstance(mask, mx.array) and mask.shape[-1] != keys.shape[-2]:
             mask = mask[..., -keys.shape[-2] :]
-
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
         )
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
         return self.o_proj(output)
@@ -117,6 +117,16 @@ class MLP(nn.Module):
         return self.down_proj(nn.gelu_approx(self.gate_proj(x)) * self.up_proj(x))
 
 
+@partial(mx.compile, shapeless=True)
+def clip_residual(x, y):
+    if x.dtype != mx.float16:
+        return x + y
+    bound = mx.finfo(mx.float16).max
+    return mx.clip(x.astype(mx.float32) + y.astype(mx.float32), -bound, bound).astype(
+        mx.float16
+    )
+
+
 class TransformerBlock(nn.Module):
     def __init__(self, args: ModelArgs, layer_idx: int):
         super().__init__()
@@ -140,9 +150,9 @@ class TransformerBlock(nn.Module):
         cache: Optional[Any] = None,
     ) -> mx.array:
         r = self.self_attn(self.input_layernorm(x), mask, cache)
-        h = x + self.post_attention_layernorm(r)
+        h = clip_residual(x, self.post_attention_layernorm(r))
         r = self.mlp(self.pre_feedforward_layernorm(h))
-        out = h + self.post_feedforward_layernorm(r)
+        out = clip_residual(h, self.post_feedforward_layernorm(r))
         return out
 
 
@@ -165,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:
@@ -202,21 +215,26 @@ class Model(nn.Module):
         self.model_type = args.model_type
         self.model = Gemma3Model(args)
         self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+        self.tie_word_embeddings = False
 
     def __call__(
         self,
         inputs: mx.array,
         cache=None,
         mask: Optional[mx.array] = None,
+        input_embeddings: Optional[mx.array] = None,
     ):
-        out = self.model(inputs, mask, cache)
-        out = self.lm_head(out)
+        out = self.model(inputs, mask, cache, input_embeddings)
+        if self.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
         return out
 
     def sanitize(self, weights):
-        weights = dict(weights)
         if "lm_head.weight" not in weights:
-            weights["lm_head.weight"] = weights["model.embed_tokens.weight"]
+            self.tie_word_embeddings = True
+            self.pop("lm_head")
         return weights
 
     @property

mlx_lm/models/gemma3n.py

@@ -0,0 +1,624 @@
+# 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
+    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[layer_idx]
+            if isinstance(config.intermediate_size, list)
+            else config.intermediate_size
+        )
+        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+        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 = self.make_cache()
+
+        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()

mlx_lm/models/glm4.py

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

mlx_lm/models/glm4_moe.py

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

mlx_lm/models/gpt2.py

@@ -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
 
@@ -133,14 +132,15 @@ class GPT2Model(nn.Module):
 
         hidden_states = self.wte(inputs)
 
-        mask = None
-        if hidden_states.shape[1] > 1:
+        offset = 0
+        if cache is not None and len(cache) > 0 and cache[0] is not None:
+            offset = cache[0].offset
 
-            position_ids = mx.array(np.arange(L))
-            hidden_states += self.wpe(position_ids)
+        position_ids = mx.arange(offset, offset + L)
+        hidden_states += self.wpe(position_ids)
 
-            if mask is None:
-                mask = create_attention_mask(hidden_states, cache)
+        if mask is None:
+            mask = create_attention_mask(hidden_states, cache)
 
         if cache is None:
             cache = [None] * len(self.h)

mlx_lm/models/gpt_bigcode.py

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

mlx_lm/models/gpt_neox.py

@@ -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
 

mlx_lm/models/gpt_oss.py

@@ -0,0 +1,379 @@
+# Copyright © 2025 Apple Inc.
+
+import math
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_causal_mask, scaled_dot_product_attention
+from .cache import KVCache, RotatingKVCache
+from .rope_utils import initialize_rope
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str = "gpt_oss"
+    num_hidden_layers: int = 36
+    num_local_experts: int = 128
+    num_experts_per_tok: int = 4
+    vocab_size: int = 201088
+    rms_norm_eps: float = 1e-05
+    hidden_size: int = 2880
+    intermediate_size: int = 2880
+    head_dim: int = 64
+    num_attention_heads: int = 64
+    num_key_value_heads: int = 8
+    sliding_window: int = 128
+    rope_theta: int = 150000
+    rope_scaling: Any = None
+    layer_types: list = None
+
+
+# These operators emulate particular methods in torch that don't exist in MLX natively
+def mlx_topk(a, k, axis=-1):
+    """MLX equivalent of torch.topk"""
+    partitioned_indices = mx.argpartition(a, kth=-k, axis=axis)
+    # Extract only the top k indices (last k elements after partition)
+    top_k_indices = partitioned_indices[..., -k:]
+    # Get the corresponding values
+    top_k_values = mx.take_along_axis(a, top_k_indices, axis=axis)
+    return top_k_values, top_k_indices
+
+
+@partial(mx.compile, shapeless=True)
+def swiglu(x_linear, x_glu, alpha: float = 1.702, limit: float = 7.0):
+    # Clamp the input values
+    x_glu = mx.clip(x_glu, a_min=None, a_max=limit)
+    x_linear = mx.clip(x_linear, a_min=-limit, a_max=limit)
+
+    glu_scaled = alpha * x_glu
+    sig = mx.sigmoid(glu_scaled)
+
+    out_glu = x_glu * sig
+    # Note we add an extra bias of 1 to the linear layer
+    return out_glu * (x_linear + 1)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def __call__(self, x, gate):
+        return swiglu(x, gate)
+
+
+class AttentionBlock(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+
+        self.head_dim = config.head_dim
+        self.num_attention_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = (
+            config.num_attention_heads // config.num_key_value_heads
+        )
+
+        self.sinks = mx.zeros((config.num_attention_heads,))
+
+        self.q_proj = nn.Linear(
+            config.hidden_size, config.num_attention_heads * self.head_dim, bias=True
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=True
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=True
+        )
+
+        self.o_proj = nn.Linear(
+            self.head_dim * config.num_attention_heads, config.hidden_size, bias=True
+        )
+
+        self.sm_scale = 1 / math.sqrt(config.head_dim)
+
+        self.rope = initialize_rope(
+            self.head_dim,
+            config.rope_theta,
+            traditional=False,
+            scaling_config=config.rope_scaling,
+        )
+
+        # Cache the mask so we don't have to create it every time
+        self._previous_mask = None
+
+    def get_causal_mask(self, x, cache):
+        _, L, _ = x.shape
+        offset = cache.offset if cache is not None else 0
+        offset = max(1, offset)
+
+        def _make_mask(L, offset):
+            zero = mx.array(0, dtype=x.dtype)
+            neginf = mx.array(-mx.inf, dtype=x.dtype)
+            mask = mx.where(create_causal_mask(L, offset - 1), zero, neginf)
+            mask = mask.reshape(1, 1, L, -1)
+            mask = mx.tile(mask, (1, self.num_attention_heads, 1, 1))
+            sinks = mx.tile(self.sinks.reshape(1, -1, 1, 1), (1, 1, L, 1))
+            mask = mx.concatenate([sinks, mask], axis=-1)
+            return mask
+
+        # When training re-create the mask so that gradients flow to the sinks.
+        # When L is large then recreate the mask because otherwise it will take
+        # a pretty significant chunk of memory.
+        if self.training or L > 8:
+            self._previous_mask = None
+            return _make_mask(L, offset)
+
+        # Create the mask once and try to reuse it. For this reason we round up
+        # to the closest multiple of 512 so we can reuse the mask several times.
+        length = ((L + offset + 511) // 512) * 512
+        if (
+            self._previous_mask is None
+            or self._previous_mask.shape[-1] < length
+            or self._previous_mask.shape[-2] != L
+        ):
+            self._previous_mask = _make_mask(L, length - L)
+
+        return self._previous_mask[..., : L + offset]
+
+    def get_sliding_window_mask(self, x, cache, window_size):
+        _, L, _ = x.shape
+        offset = cache.offset if cache is not None else 0
+        offset = max(1, offset)
+
+        def _make_mask(L, offset):
+            zero = mx.array(0, dtype=x.dtype)
+            neginf = mx.array(-mx.inf, dtype=x.dtype)
+            mask = create_causal_mask(L, offset - 1, window_size)
+            mask = mx.where(mask, zero, neginf)
+            mask = mask.reshape(1, 1, L, -1)
+            mask = mx.tile(mask, (1, self.num_attention_heads, 1, 1))
+            sinks = mx.tile(self.sinks.reshape(1, -1, 1, 1), (1, 1, L, 1))
+            mask = mx.concatenate([sinks, mask], axis=-1)
+            return mask
+
+        # If we are training then simply re-create the mask every time to make
+        # sure gradients flow to the sinks.
+        #
+        # For simplicity also re-create the mask if we have more than 1 query
+        # for now.
+        if self.training or L > 1:
+            self._previous_mask = None
+            return _make_mask(L, min(window_size + 1, offset))
+
+        # We are in inference so cache the mask and try to reuse it
+        if self._previous_mask is None:
+            self._previous_mask = _make_mask(L, window_size)
+
+        return self._previous_mask[..., : min(L + offset, window_size + 1)]
+
+    def get_mask(self, x, cache, window_size):
+        if window_size is not None:
+            return self.get_sliding_window_mask(x, cache, window_size)
+        else:
+            return self.get_causal_mask(x, cache)
+
+    def __call__(self, x: mx.array, mask: mx.array, cache=None) -> mx.array:
+        B, L, _ = x.shape
+        D = self.head_dim
+        Hk = self.num_key_value_heads
+
+        q = self.q_proj(x).reshape(B, L, -1, D).swapaxes(1, 2)
+        k = self.k_proj(x).reshape(B, L, -1, D).swapaxes(1, 2)
+        v = self.v_proj(x).reshape(B, L, -1, D).swapaxes(1, 2)
+
+        # If cache is None or the cache offset is 0 then we need to add a 0 key
+        # and value to make some space for the sink
+        if cache is None or cache.offset == 0:
+            q = self.rope(q)
+            k = self.rope(k)
+
+            zeros = mx.zeros((B, Hk, 1, D), dtype=k.dtype)
+            k = mx.concatenate([zeros, k], axis=2)
+            v = mx.concatenate([zeros, v], axis=2)
+            if cache is not None:
+                k, v = cache.update_and_fetch(k, v)
+
+        # We have already put the 0 in the cache no need to do anything special
+        else:
+            q = self.rope(q, offset=cache.offset - 1)
+            k = self.rope(k, offset=cache.offset - 1)
+            k, v = cache.update_and_fetch(k, v)
+
+        # NOTE: mask should contain the sink weights already
+        v_hat = scaled_dot_product_attention(q, k, v, cache, self.sm_scale, mask=mask)
+
+        return self.o_proj(v_hat.swapaxes(1, 2).reshape(B, L, -1))
+
+
+class MLPBlock(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+
+        self.hidden_size = config.hidden_size
+        self.num_local_experts = config.num_local_experts
+        self.num_experts_per_tok = config.num_experts_per_tok
+
+        self.experts = SwitchGLU(
+            input_dims=config.hidden_size,
+            hidden_dims=config.intermediate_size,
+            num_experts=config.num_local_experts,
+            activation=SwiGLU(),
+            bias=True,
+        )
+        self.router = nn.Linear(config.hidden_size, config.num_local_experts, bias=True)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        g = self.router(x)
+        experts, indices = mlx_topk(g, k=self.num_experts_per_tok, axis=-1)
+        expert_weights = mx.softmax(experts, axis=-1, precise=True)
+
+        # Experts block
+        x = self.experts(x, indices)
+
+        x = x * mx.expand_dims(expert_weights, axis=-1)
+        return x.sum(axis=-2)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.self_attn = AttentionBlock(config)
+        self.mlp = MLPBlock(config)
+        self.input_layernorm = nn.RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+
+    def __call__(self, x: mx.array, mask: mx.array, cache=None) -> mx.array:
+        residual = x
+        x = self.input_layernorm(x)
+        x = self.self_attn(x, mask, cache)
+        x = residual + x
+
+        residual = x
+        x = self.post_attention_layernorm(x)
+        x = self.mlp(x)
+        x = residual + x
+        return x
+
+
+class GptOssMoeModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.norm = nn.RMSNorm(args.hidden_size, args.rms_norm_eps)
+        self.layer_types = args.layer_types or [
+            "sliding_attention",
+            "full_attention",
+        ] * (args.num_hidden_layers // 2)
+        self.layers = [TransformerBlock(args) for _ in range(args.num_hidden_layers)]
+        self.window_size = args.sliding_window
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        if input_embeddings is not None:
+            x = input_embeddings
+        else:
+            x = self.embed_tokens(inputs)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        if mask is None:
+            masks = [
+                l.self_attn.get_mask(
+                    x, c, self.window_size if lt == "sliding_attention" else None
+                )
+                for (l, c, lt) in zip(self.layers, cache, self.layer_types)
+            ]
+        else:
+            masks = [mask] * len(self.layers)
+
+        for i, (layer, c, m) in enumerate(zip(self.layers, cache, masks)):
+            x = layer(x, m, c)
+        x = self.norm(x)
+        return x
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = GptOssMoeModel(args)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(self, inputs: mx.array, mask: mx.array = None, cache=None):
+        return self.lm_head(self.model(inputs, mask, cache))
+
+    def sanitize(self, weights):
+        if any("gate_proj.weight" in k for k in weights.keys()):
+            return weights  # already sanitized
+
+        new_weights = {}
+        for k, v in weights.items():
+            if "gate_up_proj" in k and "bias" not in k:
+                if "_blocks" in k:
+                    v = v.view(mx.uint32).flatten(-2)
+                    k = k.replace("_blocks", ".weight")
+                if "_scales" in k:
+                    k = k.replace("_scales", ".scales")
+                new_weights[k.replace("gate_up_proj", "gate_proj")] = mx.contiguous(
+                    v[..., ::2, :]
+                )
+                new_weights[k.replace("gate_up_proj", "up_proj")] = mx.contiguous(
+                    v[..., 1::2, :]
+                )
+            elif "down_proj" in k and "bias" not in k:
+                if "_blocks" in k:
+                    v = v.view(mx.uint32).flatten(-2)
+                    k = k.replace("_blocks", ".weight")
+                if "_scales" in k:
+                    k = k.replace("_scales", ".scales")
+                new_weights[k] = v
+            elif "gate_up_proj_bias" in k:
+                new_weights[k.replace("gate_up_proj_bias", "gate_proj.bias")] = (
+                    mx.contiguous(v[..., ::2])
+                )
+                new_weights[k.replace("gate_up_proj_bias", "up_proj.bias")] = (
+                    mx.contiguous(v[..., 1::2])
+                )
+            elif "down_proj_bias" in k:
+                new_weights[k.replace("down_proj_bias", "down_proj.bias")] = v
+            else:
+                new_weights[k] = v
+
+        return new_weights
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if path.endswith("router"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate
+
+    def make_cache(self):
+        caches = []
+        for lt in self.model.layer_types:
+            if lt == "full_attention":
+                caches.append(KVCache())
+            else:
+                caches.append(
+                    RotatingKVCache(max_size=self.args.sliding_window + 1, keep=1)
+                )
+        return caches

mlx_lm/models/granitemoe.py

@@ -0,0 +1,237 @@
+# Copyright © 2023-2024 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
+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
+    logits_scaling: float
+    attention_multiplier: float
+    embedding_multiplier: float
+    residual_multiplier: float
+    max_position_embeddings: int
+    num_key_value_heads: int
+    attention_bias: bool
+    rope_theta: float
+    num_local_experts: int
+    num_experts_per_tok: int
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = True
+
+
+class GraniteMoeAttention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+
+        self.head_dim = head_dim = args.hidden_size // n_heads
+
+        self.scale = args.attention_multiplier
+        attention_bias = args.attention_bias
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=attention_bias)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=attention_bias)
+
+        self.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            False,
+            args.rope_scaling,
+            args.max_position_embeddings,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        queries = queries.reshape(B, L, self.n_heads, -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 GraniteMoeTopKGating(nn.Module):
+    def __init__(self, input_size: int, num_experts: int, top_k: int):
+        super().__init__()
+        self.num_experts = num_experts
+        self.input_size = input_size
+        self.top_k = top_k
+        self.layer = nn.Linear(input_size, num_experts, bias=False)
+
+    def __call__(self, hidden_states: mx.array):
+        logits = self.layer(hidden_states)
+        top_k_idx = mx.argpartition(logits, kth=-self.top_k, axis=-1)[
+            ..., -self.top_k :
+        ]
+        top_k_logits = mx.take_along_axis(logits, top_k_idx, axis=-1)
+        top_k_gates = mx.softmax(top_k_logits.astype(mx.float32), axis=-1)
+        return top_k_idx, top_k_gates
+
+
+class GraniteMoeMoE(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        self.input_size = args.hidden_size
+        self.hidden_size = args.intermediate_size
+        self.switch_mlp = SwitchGLU(
+            self.input_size, self.hidden_size, args.num_local_experts
+        )
+        self.router = GraniteMoeTopKGating(
+            input_size=self.input_size,
+            num_experts=args.num_local_experts,
+            top_k=args.num_experts_per_tok,
+        )
+
+    def __call__(self, x: mx.array) -> mx.array:
+        token_ids, gates = self.router(x)
+        y = self.switch_mlp(x, token_ids)
+        return (y * gates[..., None]).sum(axis=-2).astype(y.dtype)
+
+
+class GraniteMoeDecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.self_attn = GraniteMoeAttention(args)
+        self.block_sparse_moe = GraniteMoeMoE(args)
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+        self.residual_multiplier = args.residual_multiplier
+
+    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 * self.residual_multiplier
+        r = self.block_sparse_moe(self.post_attention_layernorm(h))
+        out = h + r * self.residual_multiplier
+        return out
+
+
+class GraniteMoEModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            GraniteMoeDecoderLayer(args=args) for _ in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.embedding_multiplier = args.embedding_multiplier
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        h = self.embed_tokens(inputs) * self.embedding_multiplier
+
+        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 = GraniteMoEModel(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+        self.logits_scaling = args.logits_scaling
+
+    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 / self.logits_scaling
+
+    def sanitize(self, weights):
+        if "model.layers.0.block_sparse_moe.input_linear.weight" not in weights:
+            return weights
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}.block_sparse_moe"
+            key = f"{prefix}.input_linear.weight"
+            value = weights.pop(key)
+            gate_proj, up_proj = mx.split(value, 2, axis=1)
+            weights[key.replace("input_linear", "switch_mlp.gate_proj")] = gate_proj
+            weights[key.replace("input_linear", "switch_mlp.up_proj")] = up_proj
+            key = f"{prefix}.output_linear.weight"
+            weights[key.replace("output_linear", "switch_mlp.down_proj")] = weights.pop(
+                key
+            )
+
+        return weights
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if path.endswith("block_sparse_moe.router.layer"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/helium.py

@@ -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

mlx_lm/models/hunyuan.py

@@ -1,6 +1,5 @@
 # Copyright © 2023-2024 Apple Inc.
 
-import math
 from dataclasses import dataclass
 from typing import Any, Dict, Optional, Tuple, Union
 
@@ -30,6 +29,7 @@ class ModelArgs(BaseModelArgs):
     rope_theta: float
     use_cla: bool
     cla_share_factor: 2
+    moe_intermediate_size: Optional[Union[int, list]] = None
     rope_scaling: Optional[Dict[str, Union[float, str]]] = None
     tie_word_embeddings: bool = False
 
@@ -41,6 +41,12 @@ class ModelArgs(BaseModelArgs):
                 raise ValueError(f"rope_scaling must contain keys {required_keys}")
 
 
+def _int_or_list(arg, idx):
+    if isinstance(arg, list):
+        return arg[idx]
+    return arg
+
+
 class DynamicNTKAlphaRoPE(nn.Module):
     def __init__(
         self,
@@ -155,20 +161,29 @@ class Gate(nn.Module):
 
 
 class MoeBlock(nn.Module):
-    def __init__(self, args: ModelArgs):
+    def __init__(self, args: ModelArgs, layer_idx: int = 0):
         super().__init__()
         dim = args.hidden_size
         intermediate_size = args.intermediate_size
         self.use_shared_mlp = args.use_mixed_mlp_moe
 
         if args.use_mixed_mlp_moe:
-            self.shared_mlp = MLP(dim, intermediate_size * args.num_shared_expert)
+            num_shared = _int_or_list(args.num_shared_expert, layer_idx)
+            self.shared_mlp = MLP(dim, int(intermediate_size * num_shared))
 
         self.num_experts = num_experts = args.num_experts
-        self.top_k = args.moe_topk
+        self.top_k = _int_or_list(args.moe_topk, layer_idx)
 
         self.gate = Gate(dim, num_experts)
-        self.switch_mlp = SwitchGLU(dim, intermediate_size, num_experts)
+
+        # Use moe_intermediate_size if available, otherwise use intermediate_size
+        expert_intermediate_size = intermediate_size
+        if args.moe_intermediate_size is not None:
+            expert_intermediate_size = _int_or_list(
+                args.moe_intermediate_size, layer_idx
+            )
+
+        self.switch_mlp = SwitchGLU(dim, expert_intermediate_size, num_experts)
 
     def __call__(
         self,
@@ -182,7 +197,7 @@ class MoeBlock(nn.Module):
         scores = mx.take_along_axis(gates, inds, axis=-1)
 
         y = self.switch_mlp(x, inds)
-        y = (y * scores[..., None]).sum(axis=-2)
+        y = (y * scores[..., None].astype(mx.float32)).sum(axis=-2).astype(y.dtype)
 
         if self.use_shared_mlp:
             shared_expert_output = self.shared_mlp(x)
@@ -192,14 +207,14 @@ class MoeBlock(nn.Module):
 
 
 class DecoderLayer(nn.Module):
-    def __init__(self, args: ModelArgs, kv_proj: bool):
+    def __init__(self, args: ModelArgs, kv_proj: bool, layer_idx: int = 0):
         super().__init__()
         self.hidden_size = args.hidden_size
         self.self_attn = Attention(kv_proj, args)
         if args.num_experts == 1:
             self.mlp = MLP(args.hidden_size, args.intermediate_size)
         else:
-            self.mlp = MoeBlock(args)
+            self.mlp = MoeBlock(args, layer_idx)
 
         self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
         self.post_attention_layernorm = nn.RMSNorm(
@@ -235,6 +250,7 @@ class HunYuanModel(nn.Module):
             DecoderLayer(
                 args=args,
                 kv_proj=(not args.use_cla) or (i % args.cla_share_factor) == 0,
+                layer_idx=i,
             )
             for i in range(args.num_hidden_layers)
         ]

mlx_lm/models/hunyuan_v1_dense.py

@@ -0,0 +1,224 @@
+# Copyright © 2023-2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    rms_norm_eps: float
+    rope_theta: float = 10000
+    max_position_embeddings: int = 32768
+    attention_bias: bool = False
+    use_qk_norm: bool = True
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = False
+    head_dim: Optional[int] = None
+
+    def __post_init__(self):
+        if self.rope_scaling:
+            required_keys = {"alpha", "factor", "type"}
+            if not all(key in self.rope_scaling for key in required_keys):
+                raise ValueError(f"rope_scaling must contain keys {required_keys}")
+
+
+class DynamicNTKAlphaRoPE(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        base: float = 10000,
+        scaling_alpha: float = 1.0,
+    ):
+        super().__init__()
+        self.dims = dims
+        base = base * scaling_alpha ** (dims / (dims - 2))
+        self._freqs = base ** (mx.arange(0, self.dims, 2) / self.dims)
+
+    def __call__(self, x, offset: int = 0):
+        return mx.fast.rope(
+            x,
+            self.dims,
+            traditional=False,
+            base=None,
+            scale=1.0,
+            offset=offset,
+            freqs=self._freqs,
+        )
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+
+        head_dim = (
+            args.head_dim if args.head_dim is not None else args.hidden_size // n_heads
+        )
+        self.head_dim = head_dim
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=args.attention_bias)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.attention_bias)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.attention_bias)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=args.attention_bias)
+
+        self.use_qk_norm = args.use_qk_norm
+        if self.use_qk_norm:
+            self.query_layernorm = nn.RMSNorm(head_dim, args.rms_norm_eps)
+            self.key_layernorm = nn.RMSNorm(head_dim, args.rms_norm_eps)
+
+        scaling_alpha = 1.0
+        if args.rope_scaling and "alpha" in args.rope_scaling:
+            scaling_alpha = args.rope_scaling["alpha"]
+
+        self.rope = DynamicNTKAlphaRoPE(
+            head_dim,
+            base=args.rope_theta,
+            scaling_alpha=scaling_alpha,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        queries = queries.reshape(B, L, self.n_heads, self.head_dim).transpose(
+            0, 2, 1, 3
+        )
+        keys = keys.reshape(B, L, self.n_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, self.head_dim).transpose(
+            0, 2, 1, 3
+        )
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        if self.use_qk_norm:
+            queries = self.query_layernorm(queries)
+            keys = self.key_layernorm(keys)
+
+        if cache is not None:
+            keys, values = cache.update_and_fetch(keys, values)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        hidden_dim = args.intermediate_size
+
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class 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
+        )
+        self.args = args
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        out = h + r
+        return out
+
+
+class HunyuanV1DenseModel(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) 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 = HunyuanV1DenseModel(args)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        out = self.model(inputs, mask, cache)
+        return self.model.embed_tokens.as_linear(out)
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/internlm2.py

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

mlx_lm/models/internlm3.py

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

mlx_lm/models/kimi_vl.py

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

mlx_lm/models/lfm2-vl.py

@@ -0,0 +1,57 @@
+# 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 lfm2
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    text_config: dict
+
+    def __post_init__(self):
+        self.text_config["tie_word_embeddings"] = False
+        self.text_config["full_attn_idxs"] = [
+            i
+            for i, layer_type in enumerate(self.text_config["layer_types"])
+            if layer_type == "full_attention"
+        ]
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.language_model = lfm2.Model(lfm2.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
+
+    def make_cache(self):
+        return self.language_model.make_cache()

mlx_lm/models/lfm2.py

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

mlx_lm/models/llama.py

@@ -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:

mlx_lm/models/llama4.py

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

mlx_lm/models/longcat_flash.py

@@ -0,0 +1,383 @@
+import math
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import CacheList, KVCache
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    attention_method: str
+    zero_expert_type: str
+    hidden_size: int
+    ffn_hidden_size: int
+    moe_topk: int
+    expert_ffn_hidden_size: int
+    n_routed_experts: int
+    zero_expert_num: int
+    num_layers: int
+    vocab_size: int
+    max_position_embeddings: int
+    num_attention_heads: int
+    kv_lora_rank: int
+    q_lora_rank: int
+    qk_rope_head_dim: int
+    qk_nope_head_dim: int
+    v_head_dim: int
+    routed_scaling_factor: float
+    rms_norm_eps: float
+    rope_theta: float
+    mla_scale_q_lora: bool
+    mla_scale_kv_lora: bool
+    attention_bias: bool
+    norm_topk_prob: bool = False
+    router_bias: bool = False
+
+
+class LongcatFlashMLA(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.qk_rope_head_dim = args.qk_rope_head_dim
+        self.qk_nope_head_dim = args.qk_nope_head_dim
+        self.kv_lora_rank = args.kv_lora_rank
+        self.q_lora_rank = args.q_lora_rank
+        self.v_head_dim = args.v_head_dim
+
+        self.qk_head_dim = args.qk_nope_head_dim + args.qk_rope_head_dim
+        self.scale = self.qk_head_dim**-0.5
+
+        if self.q_lora_rank is None:
+            self.q_proj = nn.Linear(
+                args.hidden_size,
+                self.num_attention_heads * self.qk_head_dim,
+                bias=False,
+            )
+        else:
+            self.q_a_proj = nn.Linear(
+                args.hidden_size, self.q_lora_rank, bias=args.attention_bias
+            )
+            self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank)
+            self.q_b_proj = nn.Linear(
+                self.q_lora_rank,
+                self.num_attention_heads * self.qk_head_dim,
+                bias=False,
+            )
+
+        self.kv_a_proj_with_mqa = nn.Linear(
+            args.hidden_size,
+            self.kv_lora_rank + self.qk_rope_head_dim,
+            bias=args.attention_bias,
+        )
+        self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank)
+        self.kv_b_proj = nn.Linear(
+            self.kv_lora_rank,
+            self.num_attention_heads * (self.qk_nope_head_dim + args.v_head_dim),
+            bias=False,
+        )
+
+        self.o_proj = nn.Linear(
+            self.num_attention_heads * args.v_head_dim,
+            args.hidden_size,
+            bias=args.attention_bias,
+        )
+
+        if args.mla_scale_q_lora:
+            self.mla_scale_q_lora = (args.hidden_size / self.q_lora_rank) ** 0.5
+        if args.mla_scale_kv_lora:
+            self.mla_scale_kv_lora = (args.hidden_size / self.kv_lora_rank) ** 0.5
+
+        self.rope = nn.RoPE(
+            dims=self.qk_rope_head_dim, base=args.rope_theta, traditional=True
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, _ = x.shape
+
+        if self.q_lora_rank is None:
+            q_states = self.q_proj(x)
+        else:
+            q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x)))
+
+        q_states = q_states.reshape(B, L, -1, self.qk_head_dim).transpose(0, 2, 1, 3)
+
+        if self.mla_scale_q_lora is not None:
+            q_states = q_states * self.mla_scale_q_lora
+
+        q_pass, q_rot = mx.split(q_states, [self.qk_nope_head_dim], axis=-1)
+
+        compressed_kv = self.kv_a_proj_with_mqa(x)
+        k_pass, k_rot = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
+        k_pass = self.kv_a_layernorm(k_pass)
+
+        if self.mla_scale_kv_lora is not None:
+            k_pass = k_pass * self.mla_scale_kv_lora
+
+        key_shape = (B, L, -1, self.qk_nope_head_dim + self.v_head_dim)
+        k_pass = self.kv_b_proj(k_pass).reshape(*key_shape).transpose(0, 2, 1, 3)
+        k_pass, value_states = mx.split(k_pass, [self.qk_nope_head_dim], axis=-1)
+
+        k_rot = k_rot.reshape(B, 1, L, self.qk_rope_head_dim)
+
+        if cache is not None:
+            q_rot = self.rope(q_rot, cache.offset)
+            k_rot = self.rope(k_rot, cache.offset)
+        else:
+            q_rot = self.rope(q_rot)
+            k_rot = self.rope(k_rot)
+
+        k_rot = mx.broadcast_to(k_rot, (*k_pass.shape[:-1], k_rot.shape[-1]))
+
+        query_states = mx.concatenate([q_pass, q_rot], axis=-1)
+        key_states = mx.concatenate([k_pass, k_rot], axis=-1)
+
+        if cache is not None:
+            key_states, value_states = cache.update_and_fetch(key_states, value_states)
+
+        attn_output = scaled_dot_product_attention(
+            query_states,
+            key_states,
+            value_states,
+            cache=cache,
+            scale=self.scale,
+            mask=mask,
+        )
+
+        attn_output = attn_output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(attn_output)
+
+
+class LongcatFlashMLP(nn.Module):
+    def __init__(self, args: ModelArgs, is_expert: bool = False):
+        super().__init__()
+        hidden_size = args.expert_ffn_hidden_size if is_expert else args.ffn_hidden_size
+
+        self.gate_proj = nn.Linear(args.hidden_size, hidden_size, bias=False)
+        self.up_proj = nn.Linear(args.hidden_size, hidden_size, bias=False)
+        self.down_proj = nn.Linear(hidden_size, args.hidden_size, bias=False)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class LongcatFlashTopkRouter(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.config = args
+        self.top_k = args.moe_topk
+        self.n_routed_experts = args.n_routed_experts + args.zero_expert_num
+        self.routed_scaling_factor = args.routed_scaling_factor
+        self.norm_topk_prob = args.norm_topk_prob
+        self.router_bias = args.router_bias
+
+        self.classifier = nn.Linear(
+            args.hidden_size, self.n_routed_experts, bias=self.router_bias
+        )
+        self.e_score_correction_bias = mx.zeros((self.n_routed_experts,))
+
+    def __call__(self, hidden_states: mx.array) -> Tuple[mx.array, mx.array]:
+
+        dtype = hidden_states.dtype
+        router_logits = self.classifier(hidden_states)
+        scores = mx.softmax(router_logits, axis=-1)
+
+        corrected_scores = scores + self.e_score_correction_bias
+        topk_indices = mx.argpartition(corrected_scores, kth=-self.top_k, axis=-1)[
+            ..., -self.top_k :
+        ]
+        topk_weights = mx.take_along_axis(scores, topk_indices, axis=-1)
+
+        if self.norm_topk_prob:
+            denominator = mx.sum(topk_weights, axis=-1, keepdims=True) + 1e-20
+            topk_weights = topk_weights / denominator
+
+        topk_weights = topk_weights * self.routed_scaling_factor
+
+        return topk_indices, topk_weights.astype(dtype)
+
+
+class LongcatFlashMoE(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.config = args
+        self.num_experts_per_tok = args.moe_topk
+        self.n_routed_experts = args.n_routed_experts
+        self.zero_expert_num = args.zero_expert_num
+        self.zero_expert_type = args.zero_expert_type
+
+        self.switch_mlp = SwitchGLU(
+            args.hidden_size,
+            args.expert_ffn_hidden_size,
+            args.n_routed_experts,
+        )
+
+        self.router = LongcatFlashTopkRouter(args)
+
+    def __call__(self, hidden_states):
+
+        topk_indices, topk_weights = self.router(hidden_states)
+
+        # Process all regular experts at once
+        mask = topk_indices >= self.n_routed_experts
+        topk_indices = mx.where(mask, 0, topk_indices)
+        regular_weights = mx.where(mask, 0.0, topk_weights)
+
+        regular_outputs = self.switch_mlp(hidden_states, topk_indices)
+
+        weighted_outputs = regular_outputs * topk_weights[..., None]
+
+        # Add identity expert contribution if needed
+        assert self.zero_expert_type == "identity"
+        identity_weights = mx.where(mask, topk_weights, 0.0)
+        identity_outputs = hidden_states[..., None, :] * identity_weights[..., None]
+        weighted_outputs = weighted_outputs + identity_outputs
+
+        final_output = mx.sum(weighted_outputs, axis=-2)
+        return final_output
+
+
+class LongcatFlashDecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.hidden_size = args.hidden_size
+        self.mlp = LongcatFlashMoE(args)
+
+        self.self_attn = [LongcatFlashMLA(args) for _ in range(2)]
+        self.mlps = [LongcatFlashMLP(args, False) for _ in range(2)]
+        self.input_layernorm = [
+            nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) for _ in range(2)
+        ]
+        self.post_attention_layernorm = [
+            nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps) for _ in range(2)
+        ]
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        hidden_states = x
+        shortcut_mlp_output = None
+
+        for i in range(2):
+            residual = hidden_states
+
+            hidden_states = self.input_layernorm[i](hidden_states)
+            hidden_states = self.self_attn[i](hidden_states, mask=mask, cache=cache[i])
+            hidden_states = residual + hidden_states
+
+            residual = hidden_states
+            hidden_states = self.post_attention_layernorm[i](hidden_states)
+
+            if i == 0:
+                shortcut_mlp_output = self.mlp(hidden_states)
+
+            hidden_states = self.mlps[i](hidden_states)
+            hidden_states = residual + hidden_states
+
+            if i == 1:
+                hidden_states = hidden_states + shortcut_mlp_output
+
+        return hidden_states
+
+
+class LongcatFlashModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_layers = args.num_layers
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [LongcatFlashDecoderLayer(args) for idx in range(args.num_layers)]
+        self.norm = nn.RMSNorm(args.hidden_size, args.rms_norm_eps)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        h = self.embed_tokens(x)
+
+        if mask is None:
+            mask = create_attention_mask(
+                h, [cache[0][0]] if cache is not None else None
+            )
+
+        if cache is None:
+            cache = [None] * self.num_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 = LongcatFlashModel(args)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ):
+        out = self.model(inputs, mask, cache)
+        return self.lm_head(out)
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if path.endswith("classifier"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate
+
+    @property
+    def cast_predicate(self):
+        def predicate(k):
+            return "e_score_correction_bias" not in k
+
+        return predicate
+
+    def sanitize(self, weights):
+        for l in range(self.args.num_layers):
+            prefix = f"model.layers.{l}"
+            for n, m in [("w1", "gate_proj"), ("w2", "down_proj"), ("w3", "up_proj")]:
+                for k in ["weight", "scales", "biases"]:
+                    if f"{prefix}.mlp.experts.0.{m}.{k}" in weights:
+                        to_join = [
+                            weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}")
+                            for e in range(self.args.n_routed_experts)
+                        ]
+                        weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
+
+        new_weights = {}
+        for k, v in weights.items():
+            if k.startswith("model.mtp"):
+                continue
+            new_weights[k] = v
+        return new_weights
+
+    def make_cache(self):
+        return [CacheList(KVCache(), KVCache()) for _ in self.model.layers]

mlx_lm/models/mimo.py

@@ -0,0 +1,196 @@
+# Copyright © 2023-2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    rms_norm_eps: float
+    vocab_size: int
+    num_key_value_heads: int
+    max_position_embeddings: int = 32768
+    rope_theta: float = 10000.0
+    rope_traditional: bool = False
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = False
+    num_nextn_predict_layers: int = 2
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        assert args.num_key_value_heads is not None
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+
+        head_dim = args.hidden_size // n_heads
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=True)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=True)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=True)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
+
+        self.rope = initialize_rope(
+            head_dim,
+            base=args.rope_theta,
+            traditional=args.rope_traditional,
+            scaling_config=args.rope_scaling,
+            max_position_embeddings=args.max_position_embeddings,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.hidden_size = args.hidden_size
+        self.self_attn = Attention(args)
+        self.mlp = MLP(args.hidden_size, args.intermediate_size)
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+        self.args = args
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        out = h + r
+        return out
+
+
+class MiMoModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.num_hidden_layers = args.num_hidden_layers
+        self.num_nextn_predict_layers = args.num_nextn_predict_layers
+
+        assert self.vocab_size > 0
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            TransformerBlock(args=args) for _ in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        h = self.embed_tokens(inputs)
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, c)
+
+        h = self.norm(h)
+
+        return h
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = MiMoModel(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        out = self.model(inputs, mask, cache)
+
+        if self.args.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
+
+        return out
+
+    def sanitize(self, weights):
+        if self.args.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
+
+        return {
+            k: v
+            for k, v in weights.items()
+            if "self_attn.rotary_emb.inv_freq" not in k
+            and not k.startswith("model.mtp_layers.")
+        }
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/minicpm.py

@@ -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__(

mlx_lm/models/minicpm3.py

@@ -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,

mlx_lm/models/mistral3.py

@@ -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

mlx_lm/models/mixtral.py

@@ -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

mlx_lm/models/nemotron-nas.py

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

mlx_lm/models/nemotron_h.py

@@ -0,0 +1,387 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, List, Optional, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import KVCache, MambaCache
+
+
+@dataclass()
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    max_position_embeddings: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    attention_bias: bool
+    mamba_num_heads: int
+    mamba_head_dim: int
+    mamba_proj_bias: bool
+    ssm_state_size: int
+    conv_kernel: int
+    n_groups: int
+    time_step_limit: Tuple[float, float]
+    mlp_bias: bool
+    layer_norm_epsilon: float
+    rms_norm_eps: float
+    use_bias: bool
+    use_conv_bias: bool
+    residual_in_fp32: bool
+    head_dim: Optional[int] = None
+    hybrid_override_pattern: Optional[List[str]] = None
+
+
+class MambaRMSNormGated(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = mx.ones(hidden_size)
+
+    def __call__(self, hidden_states: mx.array, gate: mx.array = None) -> mx.array:
+        if gate is not None:
+            hidden_states = hidden_states * nn.silu(gate)
+        return mx.fast.rms_norm(hidden_states, self.weight, self.eps)
+
+
+class NemotronHMamba2Mixer(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_heads = args.mamba_num_heads
+        self.hidden_size = args.hidden_size
+        self.ssm_state_size = args.ssm_state_size
+        self.conv_kernel_size = args.conv_kernel
+        self.intermediate_size = args.mamba_num_heads * args.mamba_head_dim
+        self.n_groups = args.n_groups
+        self.head_dim = args.mamba_head_dim
+        self.time_step_limit = args.time_step_limit
+        self.heads_per_group = self.num_heads // self.n_groups
+
+        self.conv_dim = self.intermediate_size + 2 * self.n_groups * self.ssm_state_size
+
+        self.conv1d = nn.Conv1d(
+            in_channels=self.conv_dim,
+            out_channels=self.conv_dim,
+            kernel_size=args.conv_kernel,
+            padding=0,
+            groups=self.conv_dim,
+            bias=args.use_conv_bias,
+        )
+
+        projection_size = self.intermediate_size + self.conv_dim + self.num_heads
+        self.in_proj = nn.Linear(
+            self.hidden_size, projection_size, bias=args.mamba_proj_bias
+        )
+
+        self.dt_bias = mx.ones(self.num_heads)
+        self.A_log = mx.log(mx.arange(1, self.num_heads + 1, dtype=mx.float32))
+        self.D = mx.ones(self.num_heads)
+
+        self.norm = MambaRMSNormGated(
+            self.intermediate_size, eps=args.layer_norm_epsilon
+        )
+        self.out_proj = nn.Linear(
+            self.intermediate_size, self.hidden_size, bias=args.mamba_proj_bias
+        )
+
+    def _apply_conv(
+        self, conv_input: mx.array, cache: Optional[MambaCache] = None
+    ) -> mx.array:
+        if cache is not None:
+            if cache[0] is None:
+                conv_state = mx.zeros(
+                    (conv_input.shape[0], self.conv_kernel_size - 1, self.conv_dim),
+                    dtype=conv_input.dtype,
+                )
+            else:
+                conv_state = cache[0]
+            padded_input = mx.concatenate([conv_state, conv_input], axis=1)
+            cache[0] = padded_input[:, -(self.conv_kernel_size - 1) :, :]
+        else:
+            padded_input = mx.pad(
+                conv_input, [(0, 0), (self.conv_kernel_size - 1, 0), (0, 0)]
+            )
+
+        conv_output = self.conv1d(padded_input)
+        return nn.silu(conv_output)
+
+    def _ssm(
+        self,
+        hidden_states: mx.array,
+        B: mx.array,
+        C: mx.array,
+        dt: mx.array,
+        cache: Optional[MambaCache] = None,
+    ) -> mx.array:
+        batch_size, seq_len, _ = hidden_states.shape
+
+        dt = nn.softplus(dt + self.dt_bias)
+        dt = mx.clip(dt, self.time_step_limit[0], self.time_step_limit[1])
+
+        hidden_states = hidden_states.reshape(
+            batch_size, seq_len, self.num_heads, self.head_dim
+        )
+
+        B = B.reshape(batch_size, seq_len, self.n_groups, self.ssm_state_size)
+        B = mx.repeat(B, self.heads_per_group, axis=2)
+        C = C.reshape(batch_size, seq_len, self.n_groups, self.ssm_state_size)
+        C = mx.repeat(C, self.heads_per_group, axis=2)
+
+        A = -mx.exp(self.A_log.astype(mx.float32)).astype(hidden_states.dtype)
+
+        if cache is not None and cache[1] is not None:
+            h = cache[1]
+        else:
+            h = mx.zeros(
+                (batch_size, self.num_heads, self.head_dim, self.ssm_state_size),
+                dtype=hidden_states.dtype,
+            )
+
+        outputs = []
+        for t in range(seq_len):
+            dt_t = dt[:, t, :]
+            dA = mx.exp(dt_t * A)[..., None, None]
+            dB = (dt_t[..., None] * B[:, t])[..., None, :]
+
+            h = dA * h + dB * hidden_states[:, t, :, :, None]
+            y_t = (h @ C[:, t, :, :, None]).squeeze(-1) + self.D[
+                :, None
+            ] * hidden_states[:, t]
+            outputs.append(y_t)
+
+        if cache is not None:
+            cache[1] = h
+
+        y = mx.stack(outputs, axis=1)
+        return y.reshape(batch_size, seq_len, self.intermediate_size)
+
+    def __call__(
+        self,
+        hidden_states: mx.array,
+        cache: Optional[MambaCache] = None,
+    ) -> mx.array:
+
+        projected = self.in_proj(hidden_states)
+
+        gate, conv_input, dt = mx.split(
+            projected,
+            [self.intermediate_size, self.intermediate_size + self.conv_dim],
+            axis=-1,
+        )
+
+        conv_output = self._apply_conv(conv_input, cache)
+
+        hidden_states_ssm, B, C = mx.split(
+            conv_output,
+            [
+                self.intermediate_size,
+                self.intermediate_size + self.n_groups * self.ssm_state_size,
+            ],
+            axis=-1,
+        )
+        y = self._ssm(hidden_states_ssm, B, C, dt, cache)
+        y = self.norm(y, gate)
+        return self.out_proj(y)
+
+
+class NemotronHAttention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.hidden_size = args.hidden_size
+        self.num_heads = args.num_attention_heads
+        self.head_dim = (
+            args.head_dim
+            if args.head_dim is not None
+            else (args.hidden_size // args.num_attention_heads)
+        )
+        self.num_key_value_heads = args.num_key_value_heads
+        self.scale = self.head_dim**-0.5
+
+        self.q_proj = nn.Linear(
+            self.hidden_size, self.num_heads * self.head_dim, bias=args.attention_bias
+        )
+        self.k_proj = nn.Linear(
+            self.hidden_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=args.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            self.hidden_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=args.attention_bias,
+        )
+        self.o_proj = nn.Linear(
+            self.num_heads * self.head_dim, self.hidden_size, bias=args.attention_bias
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[KVCache] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries = self.q_proj(x).reshape(B, L, self.num_heads, -1).transpose(0, 2, 1, 3)
+        keys = (
+            self.k_proj(x)
+            .reshape(B, L, self.num_key_value_heads, -1)
+            .transpose(0, 2, 1, 3)
+        )
+        values = (
+            self.v_proj(x)
+            .reshape(B, L, self.num_key_value_heads, -1)
+            .transpose(0, 2, 1, 3)
+        )
+
+        if cache is not None:
+            keys, values = cache.update_and_fetch(keys, values)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+@partial(mx.compile, shapeless=True)
+def relu2(x):
+    return mx.square(nn.relu(x))
+
+
+class NemotronHMLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.up_proj = nn.Linear(
+            args.hidden_size, args.intermediate_size, bias=args.mlp_bias
+        )
+        self.down_proj = nn.Linear(
+            args.intermediate_size, args.hidden_size, bias=args.mlp_bias
+        )
+
+    def __call__(self, x):
+        return self.down_proj(relu2(self.up_proj(x)))
+
+
+class NemotronHBlock(nn.Module):
+    def __init__(self, args: ModelArgs, block_type: str):
+        super().__init__()
+        self.residual_in_fp32 = args.residual_in_fp32
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+        self.block_type = block_type
+
+        if self.block_type == "M":
+            self.mixer = NemotronHMamba2Mixer(args)
+        elif self.block_type == "*":
+            self.mixer = NemotronHAttention(args)
+        elif self.block_type == "-":
+            self.mixer = NemotronHMLP(args)
+
+    def __call__(
+        self,
+        x,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ):
+        hidden_states = self.norm(x)
+        if self.block_type == "M":
+            hidden_states = self.mixer(hidden_states, cache=cache)
+        elif self.block_type == "*":
+            hidden_states = self.mixer(hidden_states, mask=mask, cache=cache)
+        else:
+            hidden_states = self.mixer(hidden_states)
+
+        return x + hidden_states
+
+
+class NemotronHModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embeddings = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            NemotronHBlock(args, block_type)
+            for block_type in args.hybrid_override_pattern
+        ]
+        self.norm_f = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.fa_idx = 0
+        for b in args.hybrid_override_pattern:
+            if b == "*":
+                break
+            elif b == "M":
+                self.fa_idx += 1
+
+    def __call__(
+        self,
+        inputs,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ):
+        hidden_states = self.embeddings(inputs)
+
+        if mask is None:
+            attn_mask = create_attention_mask(
+                hidden_states, cache[self.fa_idx : self.fa_idx + 1]
+            )
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        cache_counter = 0
+        for layer in self.layers:
+            if layer.block_type == "M" or layer.block_type == "*":
+                c = cache[cache_counter]
+                cache_counter += 1
+            else:
+                c = None
+
+            if layer.block_type == "*":
+                mask = attn_mask
+            else:
+                mask = None
+            hidden_states = layer(hidden_states, mask=mask, cache=c)
+
+        return self.norm_f(hidden_states)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.backbone = NemotronHModel(args)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ):
+        out = self.backbone(inputs, mask=mask, cache=cache)
+        return self.lm_head(out)
+
+    @property
+    def layers(self):
+        return self.backbone.layers
+
+    def make_cache(self):
+        caches = []
+        for l in self.layers:
+            if l.block_type == "M":
+                caches.append(MambaCache())
+            elif l.block_type == "*":
+                caches.append(KVCache())
+        return caches
+
+    def sanitize(self, weights):
+        for k, v in weights.items():
+            if "conv1d.weight" in k and v.shape[-1] != 1:
+                weights[k] = v.moveaxis(2, 1)
+        return weights

mlx_lm/models/olmo.py

@@ -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

mlx_lm/models/openelm.py

@@ -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

mlx_lm/models/phi.py

@@ -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):

mlx_lm/models/phi3.py

@@ -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,

mlx_lm/models/phimoe.py

@@ -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,

mlx_lm/models/pixtral.py

@@ -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

mlx_lm/models/plamo2.py

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

mlx_lm/models/qwen2.py

@@ -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:

mlx_lm/models/qwen2_moe.py

@@ -1,6 +1,5 @@
 # Copyright © 2023-2024 Apple Inc.
 
-import math
 from dataclasses import dataclass
 from typing import Any, Dict, Optional, Union
 

mlx_lm/models/qwen2_vl.py

@@ -0,0 +1,59 @@
+# 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 qwen2
+from .base import BaseModelArgs
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    text_config: dict
+
+    @classmethod
+    def from_dict(cls, params):
+        if "text_config" not in params:
+            return cls(model_type=params["model_type"], text_config=params)
+        return cls(**params)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.language_model = qwen2.Model(qwen2.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("visual", None)
+        weights.pop("vision_tower", None)
+        weights = dict(tree_flatten(weights))
+
+        sanitized = {}
+        for key, value in weights.items():
+            if not key.startswith("language_model."):
+                key = "language_model." + key
+            sanitized[key] = value
+        return sanitized
+
+    @property
+    def layers(self):
+        return self.language_model.model.layers

mlx_lm/models/qwen3.py

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

mlx_lm/models/qwen3_moe.py

@@ -0,0 +1,249 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    num_experts: int
+    num_experts_per_tok: int
+    decoder_sparse_step: int
+    mlp_only_layers: List[int]
+    moe_intermediate_size: int
+    rms_norm_eps: float
+    vocab_size: int
+    num_key_value_heads: int
+    head_dim: int
+    rope_theta: float
+    tie_word_embeddings: bool
+    max_position_embeddings: int
+    norm_topk_prob: bool
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        assert args.num_key_value_heads is not None
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+
+        head_dim = getattr(
+            args, "head_dim", args.hidden_size // args.num_attention_heads
+        )
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
+
+        self.q_norm = nn.RMSNorm(head_dim, eps=args.rms_norm_eps)
+        self.k_norm = nn.RMSNorm(head_dim, eps=args.rms_norm_eps)
+
+        self.rope = nn.RoPE(
+            head_dim,
+            traditional=False,
+            base=args.rope_theta,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        # Prepare the queries, keys and values for the attention computation
+        queries = self.q_norm(queries.reshape(B, L, self.n_heads, -1)).transpose(
+            0, 2, 1, 3
+        )
+        keys = self.k_norm(keys.reshape(B, L, self.n_kv_heads, -1)).transpose(
+            0, 2, 1, 3
+        )
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class Qwen3MoeSparseMoeBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        dim = args.hidden_size
+        intermediate_size = args.moe_intermediate_size
+
+        self.num_experts = num_experts = args.num_experts
+        self.top_k = args.num_experts_per_tok
+        self.norm_topk_prob = args.norm_topk_prob
+
+        self.gate = nn.Linear(dim, num_experts, bias=False)
+        self.switch_mlp = SwitchGLU(dim, intermediate_size, num_experts)
+
+    def __call__(
+        self,
+        x: mx.array,
+    ):
+        gates = self.gate(x)
+        gates = mx.softmax(gates, axis=-1, precise=True)
+
+        k = self.top_k
+        inds = mx.stop_gradient(mx.argpartition(-gates, kth=k - 1, axis=-1)[..., :k])
+        scores = mx.take_along_axis(gates, inds, axis=-1)
+        if self.norm_topk_prob:
+            scores /= mx.sum(scores, axis=-1, keepdims=True)
+
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2)
+
+        return y
+
+
+class Qwen3MoeDecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.hidden_size = args.hidden_size
+        self.self_attn = Attention(args, layer_idx)
+
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+        self.args = args
+
+        if (layer_idx not in args.mlp_only_layers) and (
+            args.num_experts > 0 and (layer_idx + 1) % args.decoder_sparse_step == 0
+        ):
+            self.mlp = Qwen3MoeSparseMoeBlock(args)
+        else:
+            self.mlp = MLP(args.hidden_size, args.intermediate_size)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        out = h + r
+        return out
+
+
+class Qwen3MoeModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.num_hidden_layers = args.num_hidden_layers
+        assert self.vocab_size > 0
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            Qwen3MoeDecoderLayer(args=args, layer_idx=i)
+            for i in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        h = self.embed_tokens(inputs)
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = Qwen3MoeModel(args)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        out = self.model(inputs, mask, cache)
+        return self.lm_head(out)
+
+    def sanitize(self, weights):
+        if "model.layers.0.mlp.experts.0.up_proj.weight" not in weights:
+            return weights
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}"
+            for n in ["up_proj", "down_proj", "gate_proj"]:
+                if f"{prefix}.mlp.experts.0.{n}.weight" in weights:
+                    to_join = [
+                        weights.pop(f"{prefix}.mlp.experts.{e}.{n}.weight")
+                        for e in range(self.args.num_experts)
+                    ]
+                    weights[f"{prefix}.mlp.switch_mlp.{n}.weight"] = mx.stack(to_join)
+        return weights
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if path.endswith("mlp.gate"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/recurrent_gemma.py

@@ -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

mlx_lm/models/rope_utils.py

@@ -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__(
@@ -177,8 +236,26 @@ def initialize_rope(
             dims=dims,
             max_position_embeddings=max_position_embeddings,
             traditional=traditional,
+            scaling_factor=scaling_factor,
             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"],
+        )
+    elif rope_type == "mrope":
+        mrope_section = scaling_config.get("mrope_section", [])
+        assert (
+            len(mrope_section) == 3
+        ), f"MRoPE currently only supports 3 sections, got {len(mrope_section)}."
+        return nn.RoPE(dims, traditional=traditional, base=base)
     else:
         raise ValueError(f"Unsupported RoPE type {rope_type}")

mlx_lm/models/seed_oss.py

@@ -0,0 +1,186 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    rms_norm_eps: float
+    vocab_size: int
+    num_key_value_heads: int
+    head_dim: int
+    max_position_embeddings: Optional[int] = None
+    attention_bias: bool = False
+    attention_out_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
+
+        self.scale = head_dim**-0.5
+
+        input_bias = args.attention_bias
+        output_bias = args.attention_out_bias
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=input_bias)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=input_bias)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=input_bias)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=output_bias)
+
+        self.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            args.rope_traditional,
+            args.rope_scaling,
+            args.max_position_embeddings,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        # 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)
+        return self.o_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, dim, hidden_dim, bias=False):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=bias)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=bias)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=bias)
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.hidden_size = args.hidden_size
+        self.self_attn = Attention(args)
+        self.mlp = MLP(args.hidden_size, args.intermediate_size, args.mlp_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))
+        out = h + r
+        return out
+
+
+class SeedModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        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 = SeedModel(args)
+        self.tie_word_embeddings = args.tie_word_embeddings
+        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,
+    ):
+        h = self.model(inputs, mask=mask, cache=cache)
+        if self.tie_word_embeddings:
+            return h @ self.model.embed_tokens.weight.T
+        else:
+            return self.lm_head(h)
+
+    def sanitize(self, weights):
+        if self.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
+
+        return weights
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/smollm3.py

@@ -0,0 +1,76 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from . import llama
+
+
+@dataclass
+class ModelArgs(llama.ModelArgs):
+    model_type: str
+    no_rope_layer_interval: int = 4
+    no_rope_layers: Optional[list[int]] = None
+
+    def __post_init__(self):
+        super().__post_init__()
+        if self.no_rope_layers is None:
+            self.no_rope_layers = [
+                int((i + 1) % self.no_rope_layer_interval != 0)
+                for i in range(self.num_hidden_layers)
+            ]
+        elif len(self.no_rope_layers) != self.num_hidden_layers:
+            raise ValueError("`no_rope_layers` length mismatch")
+
+
+class NoPE(nn.Module):
+    """No-op used to disable rotary embeddings in selected layers."""
+
+    def __call__(self, x, offset: int = 0):
+        return x
+
+
+class Model(nn.Module):
+    """Wrapper around Llama that respects NoPE layers in SmolLM-3."""
+
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type: str = args.model_type
+
+        self.model = llama.LlamaModel(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+        for idx, use_rope in enumerate(args.no_rope_layers):
+            if not use_rope:
+                self.model.layers[idx].self_attn.rope = NoPE()
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: Optional[mx.array] = None,
+        cache=None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        out = self.model(inputs, mask, cache, input_embeddings)
+        if self.args.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
+        return out
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def sanitize(self, weights: dict):
+        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

mlx_lm/models/su_rope.py

@@ -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,
-        )

mlx_lm/models/switch_layers.py

@@ -1,11 +1,27 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import math
+from functools import partial
 
 import mlx.core as mx
 import mlx.nn as nn
 
 
+def _gather_sort(x, indices):
+    *_, M = indices.shape
+    indices = indices.flatten()
+    order = mx.argsort(indices)
+    inv_order = mx.argsort(order)
+    return x.flatten(0, -3)[order // M], indices[order], inv_order
+
+
+def _scatter_unsort(x, inv_order, shape=None):
+    x = x[inv_order]
+    if shape is not None:
+        x = mx.unflatten(x, 0, shape)
+    return x
+
+
 class QuantizedSwitchLinear(nn.Module):
     def __init__(
         self,
@@ -15,11 +31,12 @@ class QuantizedSwitchLinear(nn.Module):
         bias: bool = True,
         group_size: int = 64,
         bits: int = 4,
+        mode: str = "affine",
     ):
         super().__init__()
 
         scale = math.sqrt(1 / input_dims)
-        self.weight, self.scales, self.biases = mx.quantize(
+        self.weight, self.scales, *biases = mx.quantize(
             mx.random.uniform(
                 low=-scale,
                 high=scale,
@@ -27,23 +44,20 @@ class QuantizedSwitchLinear(nn.Module):
             ),
             group_size=group_size,
             bits=bits,
+            mode=mode,
         )
+        self.biases = biases[0] if biases else None
 
         if bias:
             self.bias = mx.zeros((num_experts, output_dims))
 
         self.group_size = group_size
         self.bits = bits
+        self.mode = mode
 
         # 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
@@ -56,16 +70,18 @@ class QuantizedSwitchLinear(nn.Module):
     def num_experts(self):
         return self.weight.shape[0]
 
-    def __call__(self, x, indices):
+    def __call__(self, x, indices, sorted_indices=False):
         x = mx.gather_qmm(
             x,
             self["weight"],
             self["scales"],
-            self["biases"],
+            self.get("biases"),
             rhs_indices=indices,
             transpose=True,
             group_size=self.group_size,
             bits=self.bits,
+            mode=self.mode,
+            sorted_indices=sorted_indices,
         )
         if "bias" in self:
             x = x + mx.expand_dims(self["bias"][indices], -2)
@@ -99,30 +115,58 @@ class SwitchLinear(nn.Module):
     def num_experts(self):
         return self.weight.shape[0]
 
-    def __call__(self, x, indices):
-        x = mx.gather_mm(x, self["weight"].swapaxes(-1, -2), rhs_indices=indices)
+    def __call__(self, x, indices, sorted_indices=False):
+        x = mx.gather_mm(
+            x,
+            self["weight"].swapaxes(-1, -2),
+            rhs_indices=indices,
+            sorted_indices=sorted_indices,
+        )
         if "bias" in self:
             x = x + mx.expand_dims(self["bias"][indices], -2)
         return x
 
-    def to_quantized(self, group_size: int = 64, bits: int = 4):
+    def to_quantized(self, group_size: int = 64, bits: int = 4, mode: str = "affine"):
         num_experts, output_dims, input_dims = self.weight.shape
         ql = QuantizedSwitchLinear(
-            input_dims, output_dims, num_experts, False, group_size, bits
+            input_dims,
+            output_dims,
+            num_experts,
+            False,
+            group_size,
+            bits,
+            mode=mode,
         )
-        ql.weight, ql.scales, ql.biases = mx.quantize(self.weight, group_size, bits)
+        ql.weight, ql.scales, *biases = mx.quantize(
+            self.weight, group_size, bits, mode=mode
+        )
+        ql.biases = biases[0] if biases else None
+
         if "bias" in self:
             ql.bias = self.bias
         return ql
 
 
+@partial(mx.compile, shapeless=True)
+def swiglu(x, gate):
+    return nn.silu(gate) * x
+
+
+class SwiGLU(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def __call__(self, x, gate):
+        return swiglu(x, gate)
+
+
 class SwitchGLU(nn.Module):
     def __init__(
         self,
         input_dims: int,
         hidden_dims: int,
         num_experts: int,
-        activation=nn.silu,
+        activation=SwiGLU(),
         bias: bool = False,
     ):
         super().__init__()
@@ -135,9 +179,25 @@ class SwitchGLU(nn.Module):
     def __call__(self, x, indices) -> mx.array:
         x = mx.expand_dims(x, (-2, -3))
 
-        x_up = self.up_proj(x, indices)
-        x_gate = self.gate_proj(x, indices)
-        x = self.down_proj(self.activation(x_gate) * x_up, indices)
+        # When we have many tokens, then sort them to make sure that the access
+        # of different experts is in order.
+        do_sort = indices.size >= 64
+        idx = indices
+        inv_order = None
+        if do_sort:
+            x, idx, inv_order = _gather_sort(x, indices)
+        if self.training:
+            idx = mx.stop_gradient(idx)
+        x_up = self.up_proj(x, idx, sorted_indices=do_sort)
+        x_gate = self.gate_proj(x, idx, sorted_indices=do_sort)
+        x = self.down_proj(
+            self.activation(x_up, x_gate),
+            idx,
+            sorted_indices=do_sort,
+        )
+
+        if do_sort:
+            x = _scatter_unsort(x, inv_order, indices.shape)
 
         return x.squeeze(-2)
 
@@ -148,7 +208,7 @@ class SwitchMLP(nn.Module):
         input_dims: int,
         hidden_dims: int,
         num_experts: int,
-        activation=nn.gelu_approx,
+        activation=nn.GELU(approx="precise"),
         bias: bool = False,
     ):
         super().__init__()
@@ -160,8 +220,20 @@ class SwitchMLP(nn.Module):
     def __call__(self, x, indices) -> mx.array:
         x = mx.expand_dims(x, (-2, -3))
 
-        x = self.fc1(x, indices)
+        # When we have many tokens, then sort them to make sure that the access
+        # of different experts is in order.
+        do_sort = indices.size >= 64
+        idx = indices
+        inv_order = None
+        if do_sort:
+            x, idx, inv_order = _gather_sort(x, indices)
+        if self.training:
+            idx = mx.stop_gradient(idx)
+        x = self.fc1(x, idx, sorted_indices=do_sort)
         x = self.activation(x)
-        x = self.fc2(x, indices)
+        x = self.fc2(x, idx, sorted_indices=do_sort)
+
+        if do_sort:
+            x = _scatter_unsort(x, inv_order, indices.shape)
 
         return x.squeeze(-2)

mlx_lm/perplexity.py

@@ -0,0 +1,187 @@
+# Copyright © 2025 Apple Inc.
+"""
+Evaluate perplexity (PPL) of MLX models.
+"""
+
+import argparse
+import math
+import time
+import types
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+
+from mlx_lm.tuner.datasets import load_dataset
+from mlx_lm.tuner.utils import get_total_parameters
+from mlx_lm.utils import load
+
+
+def load_data(
+    tokenizer,
+    data_path: str,
+    num_samples: int,
+    sequence_length: int,
+):
+    args = types.SimpleNamespace(
+        hf_dataset={
+            "path": data_path,
+            "train_split": "train",
+            "valid_split": "train[:1]",
+        },
+        train=True,
+        test=False,
+    )
+    dataset = load_dataset(args, tokenizer)[0]
+
+    perm = np.random.permutation(len(dataset)).tolist()
+
+    num_tokens = sequence_length * num_samples if num_samples > 0 else float("inf")
+    data = []
+    i = 0
+    while len(data) < num_tokens:
+        tokens, _ = dataset.process(dataset[perm[i]])
+        i += 1
+        data.extend(tokens)
+
+    data = mx.array(data[: (len(data) // sequence_length) * sequence_length])
+    data = data.reshape(-1, sequence_length)
+    if num_samples > 0:
+        data = data[:num_samples]
+    return data
+
+
+def eval_ppl(model, data, batch_size=8):
+    """
+    Evaluate perplexity on a dataset with standard error calculation.
+
+    Args:
+        model: The model to evaluate
+        data: Tokenized data tensor
+        batch_size: Batch size for evaluation
+
+    Returns:
+        tuple: (perplexity, standard_error)
+    """
+    all_losses = []
+
+    num_batches = (len(data) + batch_size - 1) // batch_size
+    for i, s in enumerate(range(0, len(data), batch_size)):
+        batch = data[s : s + batch_size]
+        # Forward pass: get logits for all tokens except last
+        logits = model(batch[:, :-1]).astype(mx.float32)
+
+        # Calculate cross-entropy loss with next tokens
+        losses = nn.losses.cross_entropy(logits, batch[:, 1:], reduction="none")
+        mx.eval(losses)
+        # Store individual token losses
+        all_losses.append(losses.flatten())
+
+        # Progress indicator
+        if (i + 1) % 1 == 0 or (i + 1) == num_batches:
+            print(f"  Processed {i + 1}/{num_batches} batches...", end="\r")
+
+    print()  # New line after progress
+
+    # Concatenate all losses into a single array
+    all_losses = mx.concatenate(all_losses)
+
+    # Calculate mean loss and perplexity
+    mean_loss = all_losses.mean().item()
+    ppl = math.exp(mean_loss)
+    # Calculate standard error
+    std_dev = mx.sqrt(mx.var(all_losses, ddof=1)).item()
+    num_tokens = all_losses.size
+    standard_error = std_dev / math.sqrt(num_tokens)
+    # Delta approximation for standard error of perplexity
+    standard_error_ppl = ppl * standard_error
+
+    return ppl, standard_error_ppl
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Evaluate perplexity of MLX models")
+    parser.add_argument(
+        "--model",
+        type=str,
+        required=True,
+        help="Path to model or Hugging Face model ID",
+    )
+    parser.add_argument(
+        "--batch-size", type=int, default=8, help="Batch size for evaluation"
+    )
+    parser.add_argument(
+        "--sequence-length",
+        type=int,
+        default=512,
+        help="Sequence length for evaluation",
+    )
+    parser.add_argument(
+        "--num-samples",
+        type=int,
+        default=256,
+        help="Number of samples to use (-1 for all available)",
+    )
+    parser.add_argument(
+        "--data-path",
+        type=str,
+        default="allenai/tulu-3-sft-mixture",
+        help="A Hugging Face dataset which is compatible with an mlx-lm dataset format.",
+    )
+    parser.add_argument(
+        "--seed", type=int, default=123, help="Random seed for data sampling"
+    )
+
+    args = parser.parse_args()
+
+    # Set random seed
+    np.random.seed(args.seed)
+    mx.random.seed(args.seed)
+
+    # Load model
+    print(f"Loading model from {args.model}...")
+    model, tokenizer = load(args.model)
+
+    # Count parameters
+    total_params = get_total_parameters(model)
+    print(f"Model loaded: {total_params/1e6:.1f}M parameters")
+
+    # Load evaluation data
+    print(f"\nLoading dataset...")
+    print(f"  Sequence length: {args.sequence_length}")
+
+    data = load_data(
+        tokenizer,
+        args.data_path,
+        num_samples=args.num_samples,
+        sequence_length=args.sequence_length,
+    )
+
+    print(f"  Loaded {len(data)} samples")
+
+    # Evaluate perplexity
+    print(f"\nEvaluating perplexity with batch size {args.batch_size}...")
+    start_time = time.time()
+
+    ppl, se = eval_ppl(model, data, batch_size=args.batch_size)
+
+    eval_time = time.time() - start_time
+    tokens_evaluated = data.shape[0] * (data.shape[1] - 1)  # B * (L - 1)
+    # Print results
+    print("\n" + "=" * 60)
+    print("EVALUATION RESULTS")
+    print("=" * 60)
+    print(f"Model: {args.model}")
+    print(f"Perplexity: {ppl:.3f} ± {se:.3f}")
+    print(f"Evaluation time: {eval_time:.2f} seconds")
+    print(f"Peak memory: {mx.get_peak_memory() / 1e9:.2f} GB")
+    print(f"Tokens per second: {tokens_evaluated / eval_time:.0f}")
+
+    # Additional statistics
+    print(f"\nDataset statistics:")
+    print(f"  Total samples: {len(data)}")
+    print(f"  Total tokens: {data.size}")
+
+
+if __name__ == "__main__":
+    main()

mlx_lm/quant/awq.py

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

mlx_lm/quant/dwq.py

@@ -0,0 +1,345 @@
+# Copyright © 2025 Apple Inc.
+
+import argparse
+import copy
+import time
+import types
+
+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.tuner.datasets import load_dataset
+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 (
+    fetch_from_hub,
+    get_model_path,
+    quantize_model,
+    save,
+)
+
+
+class Catcher(nn.Module):
+    def __init__(self, module):
+        super().__init__()
+        self.module = module
+
+    def __call__(self, *args, **kwargs):
+        outputs = self.module(*args, **kwargs)
+        self.outputs = outputs[0] if isinstance(outputs, tuple) else outputs
+        return outputs
+
+    def __getattr__(self, key: str):
+        if (value := self.get(key)) is not None:
+            return value
+        elif (m := self.get("module")) is not None:
+            return getattr(m, key)
+        else:
+            super(nn.Module, self).__getattribute__(key)
+
+
+def dwq_quantize(
+    model,
+    q_model,
+    opt,
+    train_data,
+    valid_data,
+    batch_size: int = 2,
+    max_seq_length: int = 2048,
+    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()
+    rank = group.rank()
+
+    def rprint(*args, **kwargs):
+        if rank == 0:
+            tqdm.write(*args, **kwargs)
+
+    def unfreeze(_, m):
+        if hasattr(m, "bits") and hasattr(m, "group_size") and m.mode == "affine":
+            m.unfreeze(keys=["scales", "biases"], recurse=False)
+
+    q_model.train()
+    q_model.apply_to_modules(unfreeze)
+    print_trainable_parameters(q_model)
+
+    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]
+
+    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_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()
+        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)
+            ]
+        )
+        act_loss = act_loss.mean()
+        loss = kl_loss + activation_loss_weight * act_loss
+        return loss, ntoks, kl_loss, act_loss
+
+    def step(inputs, targets, extra_targets, lengths, params):
+        (loss, ntoks, *_), grads = mx.value_and_grad(loss_fn)(
+            params, inputs, targets, extra_targets, lengths
+        )
+        grads = nn.average_gradients(grads)
+        params = opt.apply_gradients(grads, params)
+        return loss, ntoks, params
+
+    def validate(params, it):
+        v_loss = 0.0
+        v_kl_loss = 0.0
+        v_act_loss = 0.0
+        v_tokens = 0
+        for it, (batch, lengths) in tqdm(
+            enumerate(iterate_batches(valid_data, batch_size, max_seq_length)),
+            total=len(valid_data) // batch_size,
+            desc="Computing validation loss",
+            leave=False,
+        ):
+            batch = batch[:, :-1]
+            targets, extra_targets = forward(model, batch)
+            mx.eval(targets, extra_targets)
+            loss, ntoks, kl_loss, act_loss = loss_fn(
+                params, batch, targets, extra_targets, lengths
+            )
+            mx.eval(loss, ntoks)
+            loss = mx.distributed.all_sum(loss, stream=mx.cpu).item() / world_size
+            kl_loss = mx.distributed.all_sum(kl_loss, stream=mx.cpu).item() / world_size
+            act_loss = (
+                mx.distributed.all_sum(act_loss, stream=mx.cpu).item() / world_size
+            )
+            ntoks = mx.distributed.all_sum(ntoks, stream=mx.cpu).item()
+            v_tokens += ntoks
+            v_loss += loss * ntoks
+            v_kl_loss += kl_loss * ntoks
+            v_act_loss += act_loss * ntoks
+        loss = v_loss / v_tokens
+        kl_loss = v_kl_loss / v_tokens
+        act_loss = v_act_loss / v_tokens
+        rprint(f"Validation: {it=}, {loss=:.3f}, {kl_loss=:.3f}, {act_loss=:.3f}")
+        return loss
+
+    # Accumulate learned weights in higher precision
+    params = tree_map(
+        lambda x: x.astype(mx.float32),
+        q_model.trainable_parameters(),
+    )
+
+    total_loss = 0.0
+    total_tokens = 0
+    tokens = 0
+
+    tic = time.time()
+
+    # Compute initial validation loss
+    initial_valid_loss = valid_loss = validate(params, it=0)
+
+    for it, (batch, lengths) in (
+        pbar := tqdm(
+            enumerate(iterate_batches(train_data, batch_size, max_seq_length)),
+            total=len(train_data) // batch_size,
+        )
+    ):
+        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
+        total_loss += loss * ntoks
+        if rank == 0:
+            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
+                rprint(
+                    f"{it=}, {avg_loss=:.4f}, {total_tokens=},"
+                    f" {toks_per_sec=:.3f}, {peak_memory_gb=:.3f}",
+                )
+                tic = time.time()
+                tokens = 0
+                total_loss = 0
+        if (it + 1) % 200 == 0:
+            valid_loss = validate(params, it=it)
+
+    valid_loss = validate(params, it=it)
+    if initial_valid_loss < valid_loss:
+        rprint(
+            f"❌❌❌\n[WARNING] Final validation loss {valid_loss:.3f} is "
+            f"worse than initial validation loss {initial_valid_loss:.3f}."
+            " Model quality will likely be degraded.\n❌❌❌"
+        )
+
+    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 load_data(
+    tokenizer,
+    data_path: str,
+    num_samples: int,
+    max_seq_length: int,
+    num_valid_samples: int = 32,
+):
+    args = types.SimpleNamespace(
+        hf_dataset={
+            "path": data_path,
+            "train_split": "train",
+            "valid_split": "train[:1]",
+        },
+        train=True,
+        test=False,
+    )
+    dataset = load_dataset(args, tokenizer)[0]
+    perm = np.random.permutation(len(dataset))
+    train_perm = perm[:num_samples].tolist()
+    valid_perm = perm[num_samples : num_samples + num_valid_samples].tolist()
+
+    def process(idx):
+        tokens, offset = dataset.process(dataset[idx])
+        return (tokens[:max_seq_length], offset)
+
+    train = [process(i) for i in train_perm]
+    valid = [process(i) for i in valid_perm]
+    return train, valid
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model",
+        "-m",
+        help="A model to distill from for DWQ. If `quantized-model` is not"
+        " given the student model will be this model quantized according"
+        " to `bits` and `group-size`.",
+        required=True,
+    )
+    parser.add_argument(
+        "--quantized-model",
+        default=None,
+        help="An already quantized model (the student model) to improve with DWQ.",
+    )
+    parser.add_argument(
+        "--mlx-path", default="mlx_model", help="Path to save the quantized model."
+    )
+    parser.add_argument(
+        "--bits",
+        type=int,
+        default=4,
+        help="Bits per weight for quantization.",
+    )
+    parser.add_argument(
+        "--group-size", type=int, default=64, help="Group size for quantization."
+    )
+    parser.add_argument(
+        "--num-samples",
+        type=int,
+        default=2048,
+        help="Number of samples to use for training.",
+    )
+    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-6)
+    parser.add_argument("--batch-size", type=int, default=4)
+    parser.add_argument(
+        "--data-path",
+        type=str,
+        default="allenai/tulu-3-sft-mixture",
+        help="A Hugging Face dataset which is compatible with an mlx-lm dataset format.",
+    )
+    parser.add_argument(
+        "--grad-checkpoint",
+        action="store_true",
+        help="Use gradient checkpointing to reduce memory use.",
+    )
+    args = parser.parse_args()
+
+    group = mx.distributed.init()
+
+    num_samples = args.num_samples
+    if num_samples % group.size() > 0:
+        num_samples += group.size() - num_samples % group.size()
+
+    np.random.seed(args.seed)
+    mx.random.seed(args.seed)
+
+    model_path, hf_repo = get_model_path(args.model, revision=None)
+    model, config, tokenizer = fetch_from_hub(
+        model_path, lazy=True, trust_remote_code=True
+    )
+
+    train_data, valid_data = load_data(
+        tokenizer, args.data_path, args.num_samples, args.max_seq_length
+    )
+
+    if args.quantized_model is not None:
+        q_model_path, _ = get_model_path(args.quantized_model, revision=None)
+        q_model, config, _ = fetch_from_hub(
+            q_model_path, lazy=True, trust_remote_code=True
+        )
+        if "quantization" not in config:
+            raise ValueError("Quantized model must already be quantized.")
+    else:
+        q_model = copy.deepcopy(model)
+        _, config = quantize_model(
+            q_model,
+            config,
+            group_size=args.group_size,
+            bits=args.bits,
+        )
+
+    opt = optimizers.Adam(learning_rate=args.learning_rate, bias_correction=True)
+    dwq_quantize(
+        model,
+        q_model,
+        opt,
+        train_data,
+        valid_data,
+        batch_size=args.batch_size,
+        max_seq_length=args.max_seq_length,
+        gradient_checkpoint=args.grad_checkpoint,
+    )
+    save(
+        args.mlx_path,
+        model_path,
+        q_model,
+        tokenizer,
+        config,
+        hf_repo=hf_repo,
+    )

mlx_lm/quant/dynamic_quant.py

@@ -0,0 +1,264 @@
+# 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_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.utils import (
+    compute_bits_per_weight,
+    fetch_from_hub,
+    get_model_path,
+    load,
+    quantize_model,
+    save,
+)
+
+
+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 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)
+
+    def compute_sensitivity(gradient, low_q_weight, original_weight):
+        n_batches = (len(data) + batch_size - 1) // batch_size
+        gradient = gradient / n_batches
+        high_q_weight = qdq(original_weight, high_bits, high_group_size)
+        param_size = original_weight.size / 1e6
+        alignment = (gradient * (low_q_weight - high_q_weight)).sum()
+        return alignment / param_size
+
+    sensitivities = tree_map(
+        compute_sensitivity,
+        grad_accum,
+        q_model.parameters(),
+        model.parameters(),
+    )
+    mx.eval(sensitivities)
+
+    sensitivities = [(k[:-7], s.item()) for k, s in tree_flatten(sensitivities)]
+
+    return sensitivities
+
+
+def estimate_threshold(
+    model,
+    sensitivities,
+    target_bpw,
+    low_bits,
+    low_group_size,
+    high_bits,
+    high_group_size,
+):
+    def predicate(p, m, high_threshold):
+        if not hasattr(m, "to_quantized"):
+            return False
+        if sensitivities[p] > high_threshold:
+            return {"bits": high_bits, "group_size": high_group_size}
+        return True
+
+    # Binary search for the threshold
+    sens_vals = list(sensitivities.values())
+    min_threshold = min(sens_vals)
+    max_threshold = max(sens_vals)
+    tolerance = 1e-3 * (max_threshold - min_threshold)
+    while (max_threshold - min_threshold) > tolerance:
+        mid = (max_threshold + min_threshold) / 2
+        class_predicate = lambda p, m: predicate(p, m, mid)
+        q_model = copy.deepcopy(model)
+        nn.quantize(
+            q_model,
+            group_size=low_group_size,
+            bits=low_bits,
+            class_predicate=class_predicate,
+        )
+        bpw = compute_bits_per_weight(q_model)
+        if bpw > target_bpw:
+            min_threshold = mid
+        else:
+            max_threshold = mid
+
+    return (max_threshold + min_threshold) / 2
+
+
+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=64)
+    parser.add_argument("--high-bits", type=int, default=5)
+    parser.add_argument("--high-group-size", type=int, default=64)
+    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}")
+
+    threshold = 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,
+    )
+
+    def quant_predicate(p, m, _):
+        if not hasattr(m, "to_quantized"):
+            return False
+        if sensitivities[p] > threshold:
+            return {"bits": args.high_bits, "group_size": args.high_group_size}
+        return True
+
+    model, config = quantize_model(
+        model,
+        config,
+        group_size=args.low_group_size,
+        bits=args.low_bits,
+        quant_predicate=quant_predicate,
+    )
+
+    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()

mlx_lm/quant/gptq.py

@@ -0,0 +1,232 @@
+# Copyright © 2025 Apple Inc.
+
+"""
+Implements GPTQ
+
+- https://arxiv.org/abs/2210.17323
+- https://github.com/AutoGPTQ
+"""
+
+import argparse
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten, tree_unflatten
+from tqdm import tqdm
+
+from mlx_lm.models.switch_layers import QuantizedSwitchLinear, SwitchLinear
+from mlx_lm.quant.utils import load_data
+from mlx_lm.utils import (
+    compute_bits_per_weight,
+    fetch_from_hub,
+    get_model_path,
+    save,
+)
+
+
+def quantize(w, bits, scales, biases):
+    assert bits in {2, 4, 8}, f"Unsupported bits {bits}"
+    el_per_int = 32 // bits
+    n_bins = 2**bits - 1
+    w = mx.unflatten(w, -1, (scales.shape[-1], -1))
+    w = mx.clip(
+        mx.round((w - biases[..., None]) / scales[..., None]), 0.0, n_bins
+    ).astype(mx.uint32)
+    shifts = mx.power(2, mx.arange(0, 32, bits, mx.uint32))
+    w = mx.unflatten(w, -1, (-1, el_per_int))
+    w = mx.sum(w * shifts, axis=-1)
+    return w.flatten(-2, -1)
+
+
+class Catcher(nn.Module):
+    def __init__(self, module):
+        super().__init__()
+        self.module = module
+        self.H = mx.array(0.0)
+
+    def __call__(self, x, *args, **kwargs):
+        xf = x.flatten(0, -2)
+        self.H = self.H + xf.T @ xf
+        return self.module(x, *args, **kwargs)
+
+
+def gptq_quantize(
+    model,
+    data,
+    bits,
+    group_size,
+    fallback_bits,
+    fallback_group_size,
+    batch_size=8,
+):
+    layers = []
+    gptq_types = {nn.Linear, SwitchLinear}
+    for k, l in tree_flatten(model.leaf_modules(), is_leaf=nn.Module.is_module):
+        if type(l) in gptq_types:
+            layers.append((k, Catcher(l)))
+    model.update_modules(tree_unflatten(layers))
+
+    # Evaluate the Hessians for all quantizable layers
+    for e, s in tqdm(
+        enumerate(range(0, len(data), batch_size)),
+        total=len(data) // batch_size,
+        desc="Computing Hessians",
+    ):
+        batch = data[s : s + batch_size]
+        model(batch)
+        mx.eval(layers)
+
+    def compute_inverse_hessian(H):
+        with mx.stream(mx.cpu):
+            damp = 1e-2 * mx.mean(mx.diag(H))
+            diag = mx.arange(H.shape[0])
+            H[diag, diag] += damp
+            H = mx.linalg.cholesky(H)
+            H = mx.linalg.cholesky_inv(H)
+            Hinv = mx.linalg.cholesky(H, upper=True)
+            return Hinv
+
+    @mx.compile
+    def gptq_error(w, d, scales, biases):
+        n_bins = 2**bits - 1
+        q = mx.clip(mx.round((w - biases) / scales), 0.0, n_bins)
+        q = scales * q + biases
+        return (w - q) / d
+
+    for lid, (key, l) in tqdm(
+        enumerate(layers),
+        total=len(layers),
+        desc="Quantizing",
+    ):
+        Hinv = compute_inverse_hessian(l.H)
+        del l.H
+        mx.eval(Hinv)
+
+        orig_type = l.module.weight.dtype
+        W = l.module.weight.astype(mx.float32)
+
+        all_scales = []
+        all_biases = []
+        for i in range(0, W.shape[-1], group_size):
+            j = i + group_size
+            Wl = W[..., i:j]
+            err = mx.zeros_like(Wl)
+
+            # Find scales and biases
+            _, scales, biases = mx.quantize(Wl, bits=bits, group_size=group_size)
+
+            all_scales.append(scales)
+            all_biases.append(biases)
+            for k in range(group_size):
+                k += i
+                w = W[..., k : k + 1]
+                d = Hinv[k, k]
+
+                e = gptq_error(w, d, scales, biases)
+
+                W[..., k : k + j] -= e @ Hinv[k : k + 1, k : k + j]
+                err[..., k : k + 1] = e
+                mx.eval(err, W)
+
+            W[..., j:] -= err @ Hinv[i:j, j:]
+
+        # Quantize with the given scales and biases
+        scales = mx.concatenate(all_scales, axis=-1)
+        biases = mx.concatenate(all_biases, axis=-1)
+        Wq = quantize(W, bits, scales, biases)
+        layer = l.module.to_quantized(bits=bits, group_size=group_size)
+        layer.weight = Wq
+        layer.scales = scales
+        layer.biases = biases
+        layer.set_dtype(orig_type)
+        mx.eval(layer)
+        layers[lid] = (key, layer)
+
+    model.update_modules(tree_unflatten(layers))
+
+    layers = tree_flatten(
+        model.leaf_modules(),
+        is_leaf=nn.Module.is_module,
+    )
+    config = {"bits": bits, "group_size": group_size}
+    fallback_config = {"bits": fallback_bits, "group_size": fallback_group_size}
+    q_layers = []
+    for e, (k, l) in enumerate(layers):
+        if hasattr(l, "to_quantized"):
+            config[k] = fallback_config
+            q_layers.append((k, l.to_quantized(**fallback_config)))
+    if len(q_layers) > 0:
+        model.update_modules(tree_unflatten(q_layers))
+    return model, 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")
+    parser.add_argument(
+        "--bits", type=int, default=4, help="Quantization bits for GPTQ layers"
+    )
+    parser.add_argument(
+        "--group-size",
+        type=int,
+        default=64,
+        help="Quantization group size for GPTQ layers",
+    )
+    parser.add_argument(
+        "--fallback-bits",
+        type=int,
+        default=6,
+        help="Quantization bits for non-GPTQ layers",
+    )
+    parser.add_argument(
+        "--fallback-group-size",
+        type=int,
+        default=64,
+        help="Quantization group size for non-GPTQ layers",
+    )
+    parser.add_argument(
+        "--num-samples",
+        type=int,
+        default=-1,
+        help="Number of samples from the calibration dataset, use -1 for all.",
+    )
+    parser.add_argument(
+        "--sequence-length",
+        type=int,
+        default=512,
+        help="Sequence length for the calibration data.",
+    )
+    parser.add_argument("--seed", type=int, default=123)
+    args = parser.parse_args()
+
+    mx.random.seed(args.seed)
+
+    model_path, hf_repo = get_model_path(args.model, revision=None)
+    model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
+    calibration_data = load_data(tokenizer, args.num_samples, args.sequence_length)
+
+    model, config["quantization"] = gptq_quantize(
+        model,
+        calibration_data,
+        args.bits,
+        args.group_size,
+        args.fallback_bits,
+        args.fallback_group_size,
+    )
+
+    bpw = compute_bits_per_weight(model)
+    print(f"Quantized model with {bpw:.3f} bits per weight.")
+
+    save(
+        args.mlx_path,
+        model_path,
+        model,
+        tokenizer,
+        config,
+        hf_repo=hf_repo,
+    )
+
+
+if __name__ == "__main__":
+    main()

mlx_lm/quant/utils.py

@@ -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]

mlx_lm/sample_utils.py

@@ -2,7 +2,7 @@
 
 import math
 from functools import partial
-from typing import Callable, Dict, Optional
+from typing import Callable, Dict, List, Optional
 
 import mlx.core as mx
 
@@ -12,8 +12,11 @@ def make_sampler(
     top_p: float = 0.0,
     min_p: float = 0.0,
     min_tokens_to_keep: int = 1,
-    top_k: int = -1,
-) -> Callable[mx.array, mx.array]:
+    top_k: int = 0,
+    xtc_probability: float = 0.0,
+    xtc_threshold: float = 0.0,
+    xtc_special_tokens: List[int] = [],
+) -> Callable[[mx.array], mx.array]:
     """
     Make a sampler function for use with ``generate_step``.
 
@@ -28,6 +31,13 @@ def make_sampler(
           be filtered by min_p sampling.
         top_k (int, optional): The top k tokens ranked by probability to constrain
           the sampling to.
+        xtc_probability (float, optional): The probability of applying XTC
+            sampling.
+        xtc_threshold (float, optional): The threshold the probs need to reach
+            for being sampled.
+        xtc_special_tokens (list(int), optional): List of special tokens IDs to
+            be excluded from XTC sampling.
+
 
     Returns:
         Callable[mx.array, mx.array]:
@@ -38,20 +48,23 @@ def make_sampler(
 
     # Create sampler chain
     sampling_methods = []
-    if top_k > 0:
-        sampling_methods.append(lambda x: apply_top_k(x, top_k))
     if top_p > 0 and top_p < 1.0:
         sampling_methods.append(lambda x: apply_top_p(x, top_p))
     if min_p != 0.0:
         sampling_methods.append(lambda x: apply_min_p(x, min_p, min_tokens_to_keep))
+    if xtc_probability > 0.0:
+        sampling_methods.append(
+            lambda x: apply_xtc(x, xtc_probability, xtc_threshold, xtc_special_tokens)
+        )
+    if top_k > 0:
+        sampling_methods.append(lambda x: apply_top_k(x, top_k))
 
     # Apply the sampling methods
-    def sampler(logits):
+    def sampler(logprobs):
         for method in sampling_methods:
-            logits = method(logits)
-
+            logprobs = method(logprobs)
         # Return the sampled token
-        return categorical_sampling(logits, temp)
+        return categorical_sampling(logprobs, temp)
 
     return sampler
 
@@ -170,8 +183,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(
@@ -182,41 +199,76 @@ 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)
 
+    # 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)
+
     # select tokens with cumulative probs below threshold
-    top_probs = mx.where(
+    return mx.where(
         cumulative_probs > 1 - top_p,
-        sorted_probs,
-        0,
+        logprobs,
+        -float("inf"),
     )
 
-    # 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)
+@partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)
+def apply_xtc(
+    logits: mx.array,
+    xtc_probability: float,
+    xtc_threshold: float,
+    xtc_special_tokens: List[int],
+) -> mx.array:
+    """
+    Apply XTC sampling to the logits.
 
-    # Convert back to logits and return
-    return mx.log(original_order_probs)
+    Args:
+        logits: The logits from the model's output.
+        xtc_probability (float): Probability of XTC sampling to happen for each token
+        xtc_threshold (float): The threshold the probs need to reach for being sampled.
+        special_tokens_ids (list(int)): List of special tokens IDs to be excluded from XTC sampling.
+    """
+    if not (0 <= xtc_threshold <= 0.5):
+        raise ValueError(
+            f"`threshold` has to be a float in the [0, 0.5] interval, but is {xtc_threshold}"
+        )
+    if not (0 <= xtc_probability <= 1.0):
+        raise ValueError(
+            f"`probability` has to be a float in the [0, 1] interval, but is {xtc_probability}"
+        )
+
+    probs = mx.softmax(logits, -1)
+    mask = probs > mx.where(probs > xtc_threshold, probs, mx.inf).min()
+    if xtc_special_tokens:
+        mask[..., xtc_special_tokens] = False
+
+    return mx.where(
+        mx.random.uniform(0, 1) > xtc_probability,
+        logits,
+        mx.where(mask, -mx.inf, logits),
+    )
 
 
 @partial(mx.compile, inputs=mx.random.state, outputs=mx.random.state)

mlx_lm/server.py

@@ -30,7 +30,7 @@ from ._version import __version__
 from .generate import stream_generate
 from .models.cache import can_trim_prompt_cache, make_prompt_cache, trim_prompt_cache
 from .sample_utils import make_logits_processors, make_sampler
-from .utils import load
+from .utils import common_prefix_len, load
 
 
 def get_system_fingerprint():
@@ -141,12 +141,14 @@ 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
 class PromptCache:
     cache: List[Any] = field(default_factory=list)
-    model_key: Tuple[str, Optional[str]] = ("", None)
+    model_key: Tuple[str, Optional[str]] = ("", None, None)
     tokens: List[int] = field(default_factory=list)
 
 
@@ -157,10 +159,13 @@ class ModelProvider:
         self.model_key = None
         self.model = None
         self.tokenizer = None
+        self.draft_model = None
 
         # Preload the default model if it is provided
+        self.default_model_map = {}
         if self.cli_args.model is not None:
-            self.load("default_model")
+            self.default_model_map[self.cli_args.model] = "default_model"
+            self.load(self.cli_args.model, draft_model_path="default_model")
 
     def _validate_model_path(self, model_path: str):
         model_path = Path(model_path)
@@ -170,14 +175,21 @@ class ModelProvider:
             )
 
     # Added in adapter_path to load dynamically
-    def load(self, model_path, adapter_path=None):
-        if self.model_key == (model_path, adapter_path):
+    def load(self, model_path, adapter_path=None, draft_model_path=None):
+        model_path, adapter_path, draft_model_path = map(
+            lambda s: s.lower() if s else None,
+            (model_path, adapter_path, draft_model_path),
+        )
+
+        model_path = self.default_model_map.get(model_path, model_path)
+        if self.model_key == (model_path, adapter_path, draft_model_path):
             return self.model, self.tokenizer
 
         # Remove the old model if it exists.
         self.model = None
         self.tokenizer = None
         self.model_key = None
+        self.draft_model = None
 
         # Building tokenizer_config
         tokenizer_config = {
@@ -186,12 +198,16 @@ class ModelProvider:
         if self.cli_args.chat_template:
             tokenizer_config["chat_template"] = self.cli_args.chat_template
 
-        if model_path == "default_model" and self.cli_args.model is not None:
+        if model_path == "default_model":
+            if self.cli_args.model is None:
+                raise ValueError(
+                    "A model path has to be given as a CLI "
+                    "argument or in the HTTP request"
+                )
+            adapter_path = adapter_path or self.cli_args.adapter_path
             model, tokenizer = load(
                 self.cli_args.model,
-                adapter_path=(
-                    adapter_path if adapter_path else self.cli_args.adapter_path
-                ),  # if the user doesn't change the model but adds an adapter path
+                adapter_path=adapter_path,
                 tokenizer_config=tokenizer_config,
             )
         else:
@@ -204,10 +220,30 @@ class ModelProvider:
             if tokenizer.chat_template is None:
                 tokenizer.chat_template = tokenizer.default_chat_template
 
-        self.model_key = (model_path, adapter_path)
+        self.model_key = (model_path, adapter_path, draft_model_path)
         self.model = model
         self.tokenizer = tokenizer
 
+        def validate_draft_tokenizer(draft_tokenizer):
+            # Check if tokenizers are compatible
+            if draft_tokenizer.vocab_size != tokenizer.vocab_size:
+                logging.warning(
+                    "Draft model tokenizer does not match model tokenizer. "
+                    "Speculative decoding may not work as expected."
+                )
+
+        # Load draft model if specified
+        if (
+            draft_model_path == "default_model"
+            and self.cli_args.draft_model is not None
+        ):
+            self.draft_model, draft_tokenizer = load(self.cli_args.draft_model)
+            validate_draft_tokenizer(draft_tokenizer)
+
+        elif draft_model_path is not None and draft_model_path != "default_model":
+            self._validate_model_path(draft_model_path)
+            self.draft_model, draft_tokenizer = load(draft_model_path)
+            validate_draft_tokenizer(draft_tokenizer)
         return self.model, self.tokenizer
 
 
@@ -268,7 +304,23 @@ class APIHandler(BaseHTTPRequestHandler):
         # Fetch and parse request body
         content_length = int(self.headers["Content-Length"])
         raw_body = self.rfile.read(content_length)
-        self.body = json.loads(raw_body.decode())
+        try:
+            self.body = json.loads(raw_body.decode())
+        except json.JSONDecodeError as e:
+            logging.error(f"JSONDecodeError: {e} - Raw body: {raw_body.decode()}")
+            # Set appropriate headers based on streaming requirement
+            if self.stream:
+                self._set_stream_headers(400)
+                self.wfile.write(
+                    f"data: {json.dumps({'error': f'Invalid JSON in request body: {e}'})}\n\n".encode()
+                )
+            else:
+                self._set_completion_headers(400)
+                self.wfile.write(
+                    json.dumps({"error": f"Invalid JSON in request body: {e}"}).encode()
+                )
+            return
+
         indent = "\t"  # Backslashes can't be inside of f-strings
         logging.debug(f"Incoming Request Body: {json.dumps(self.body, indent=indent)}")
         assert isinstance(
@@ -279,22 +331,38 @@ class APIHandler(BaseHTTPRequestHandler):
         self.stream = self.body.get("stream", False)
         self.stream_options = self.body.get("stream_options", None)
         self.requested_model = self.body.get("model", "default_model")
+        self.requested_draft_model = self.body.get("draft_model", "default_model")
+        self.num_draft_tokens = self.body.get(
+            "num_draft_tokens", 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)
+        self.xtc_threshold = self.body.get("xtc_threshold", 0.0)
         self.logit_bias = self.body.get("logit_bias", None)
         self.logprobs = self.body.get("logprobs", -1)
+        self.seed = self.body.get("seed", None)
         self.validate_model_parameters()
-
+        if self.seed is not None:
+            mx.random.seed(self.seed)
         # Load the model if needed
         try:
             self.model, self.tokenizer = self.model_provider.load(
-                self.requested_model, self.adapter
+                self.requested_model,
+                self.adapter,
+                self.requested_draft_model,
             )
         except:
             self._set_completion_headers(404)
@@ -338,6 +406,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
@@ -363,11 +440,21 @@ class APIHandler(BaseHTTPRequestHandler):
                 self.logit_bias = {int(k): v for k, v in self.logit_bias.items()}
             except ValueError:
                 raise ValueError("logit_bias must be a dict of int to float")
-
+        if not (
+            isinstance(self.xtc_probability, float)
+            and 0.00 <= self.xtc_probability <= 1.00
+        ):
+            raise ValueError(f"xtc_probability must be a float between 0.00 and 1.00")
+        if not (
+            isinstance(self.xtc_threshold, float) and 0.00 <= self.xtc_threshold <= 0.50
+        ):
+            raise ValueError(f"xtc_threshold must be a float between 0.00 and 0.5")
         if not isinstance(self.requested_model, str):
             raise ValueError("model must be a string")
         if self.adapter is not None and not isinstance(self.adapter, str):
             raise ValueError("adapter must be a string")
+        if self.seed is not None and not isinstance(self.seed, int):
+            raise ValueError("seed must be an integer")
 
     def generate_response(
         self,
@@ -378,6 +465,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
@@ -396,13 +484,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 = {
@@ -414,16 +515,18 @@ class APIHandler(BaseHTTPRequestHandler):
             "choices": [
                 {
                     "index": 0,
-                    "logprobs": {
-                        "token_logprobs": token_logprobs,
-                        "top_logprobs": top_logprobs,
-                        "tokens": tokens,
-                    },
                     "finish_reason": finish_reason,
-                }
+                },
             ],
         }
 
+        if token_logprobs or top_logprobs or tokens:
+            response["choices"][0]["logprobs"] = {
+                "token_logprobs": token_logprobs,
+                "top_logprobs": top_logprobs,
+                "tokens": tokens,
+            }
+
         if not self.stream:
             if not (
                 isinstance(prompt_token_count, int)
@@ -444,42 +547,99 @@ 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:
-            ValueError(f"Unsupported response type: {self.object_type}")
+            raise ValueError(f"Unsupported response type: {self.object_type}")
 
         return response
 
+    def reset_prompt_cache(self, prompt):
+        """Resets the prompt cache and associated state.
+
+        Args:
+            prompt (List[int]): The tokenized new prompt which will populate the
+                reset cache.
+        """
+        logging.debug(f"*** Resetting cache. ***")
+        self.prompt_cache.model_key = self.model_provider.model_key
+        self.prompt_cache.cache = make_prompt_cache(self.model_provider.model)
+        if self.model_provider.draft_model is not None:
+            self.prompt_cache.cache += make_prompt_cache(
+                self.model_provider.draft_model
+            )
+        self.prompt_cache.tokens = list(prompt)  # Cache the new prompt fully
+
     def get_prompt_cache(self, prompt):
+        """
+        Determines the portion of the prompt that needs processing by comparing
+        it to the cached prompt and attempting to reuse the common prefix.
+
+        This function updates the internal prompt cache state (tokens and model cache)
+        based on the comparison. If a common prefix exists, it attempts to trim
+        the model cache (if supported) to match the common prefix length, avoiding
+        recomputation.
+
+        Args:
+            prompt (List[int]): The tokenized new prompt.
+
+        Returns:
+            List[int]: The suffix of the prompt that actually needs to be processed
+                       by the model. This will be the full prompt if the cache is
+                       reset or cannot be effectively used.
+        """
         cache_len = len(self.prompt_cache.tokens)
         prompt_len = len(prompt)
-        prefix_len = min(cache_len, prompt_len)
+        com_prefix_len = common_prefix_len(self.prompt_cache.tokens, prompt)
+
+        # 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
-            or prompt[:prefix_len] != self.prompt_cache.tokens[:prefix_len]
+            or com_prefix_len == 0
         ):
-            self.prompt_cache.model_key = self.model_provider.model_key
-            self.prompt_cache.cache = make_prompt_cache(self.model_provider.model)
-            self.prompt_cache.tokens = []
-        elif cache_len >= prompt_len:
-            # Trim the cache if it contains the prompt as a prefix. This case
-            # is useful for reusing the cache for multiple queries with a long
-            # prompt
+            self.reset_prompt_cache(prompt)
+
+        # Condition 2: Common prefix exists and matches cache length. Process suffix.
+        elif com_prefix_len == cache_len:
+            logging.debug(
+                f"*** Cache is prefix of prompt (cache_len: {cache_len}, prompt_len: {prompt_len}). Processing suffix. ***"
+            )
+            prompt = prompt[com_prefix_len:]
+            self.prompt_cache.tokens.extend(prompt)
+
+        # Condition 3: Common prefix exists but is shorter than cache length. Attempt trim.
+        elif com_prefix_len < cache_len:
+            logging.debug(
+                f"*** Common prefix ({com_prefix_len}) shorter than cache ({cache_len}). Attempting trim. ***"
+            )
+
             if can_trim_prompt_cache(self.prompt_cache.cache):
-                num_to_trim = cache_len - prompt_len + 1
+                num_to_trim = cache_len - com_prefix_len
+                logging.debug(f"    Trimming {num_to_trim} tokens from cache.")
                 trim_prompt_cache(self.prompt_cache.cache, num_to_trim)
-                self.prompt_cache.tokens = self.prompt_cache.tokens[:-num_to_trim]
-                prompt = prompt[-1:]
+                self.prompt_cache.tokens = self.prompt_cache.tokens[:com_prefix_len]
+                prompt = prompt[com_prefix_len:]
+                self.prompt_cache.tokens.extend(prompt)
             else:
-                self.prompt_cache.cache = make_prompt_cache(self.model_provider.model)
-                self.prompt_cache.tokens = []
+                logging.debug(f"    Cache cannot be trimmed. Resetting cache.")
+                self.reset_prompt_cache(prompt)
+
+        # This case should logically not be reached if com_prefix_len <= cache_len
         else:
-            # Trim the prompt if it contains the cache as a prefix. This case
-            # is to avoid recomputing the cache in multi-turn chats.
-            prompt = prompt[cache_len:]
-        self.prompt_cache.tokens.extend(prompt)
+            logging.error(
+                f"Unexpected cache state: com_prefix_len ({com_prefix_len}) > cache_len ({cache_len}). Resetting cache."
+            )
+            self.reset_prompt_cache(prompt)
+
+        logging.debug(f"Returning {len(prompt)} tokens for processing.")
         return prompt
 
     def handle_completion(
@@ -510,10 +670,45 @@ class APIHandler(BaseHTTPRequestHandler):
 
         text = ""
         tic = time.perf_counter()
-        sampler = make_sampler(self.temperature, top_p=self.top_p)
-        logits_processors = make_logits_processors(
-            self.logit_bias, self.repetition_penalty, self.repetition_context_size
+        sampler = make_sampler(
+            self.temperature,
+            top_p=self.top_p,
+            top_k=self.top_k,
+            min_p=self.min_p,
+            xtc_probability=self.xtc_probability,
+            xtc_threshold=self.xtc_threshold,
+            xtc_special_tokens=[
+                self.tokenizer.eos_token_id,
+                self.tokenizer.encode("\n"),
+            ],
         )
+        logits_processors = make_logits_processors(
+            self.logit_bias,
+            self.repetition_penalty,
+            self.repetition_context_size,
+        )
+
+        tool_calls = []
+        tool_text = ""
+        in_tool_call = False
+        segment = ""
+
+        # Create keepalive callback to send SSE comments during long prompt processing
+        def keepalive_callback(processed_tokens, total_tokens):
+            logging.info(
+                f"Prompt processing progress: {processed_tokens}/{total_tokens}"
+            )
+            if self.stream:
+                try:
+                    # Send SSE comment for keepalive - invisible to clients but keeps connection alive
+                    self.wfile.write(
+                        f": keepalive {processed_tokens}/{total_tokens}\n\n".encode()
+                    )
+                    self.wfile.flush()
+                except (BrokenPipeError, ConnectionResetError, OSError):
+                    # Client disconnected, ignore
+                    pass
+
         for gen_response in stream_generate(
             model=self.model,
             tokenizer=self.tokenizer,
@@ -522,13 +717,31 @@ class APIHandler(BaseHTTPRequestHandler):
             sampler=sampler,
             logits_processors=logits_processors,
             prompt_cache=self.prompt_cache.cache,
+            draft_model=self.model_provider.draft_model,
+            num_draft_tokens=self.num_draft_tokens,
+            prompt_progress_callback=keepalive_callback,
         ):
-            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)
+            self.prompt_cache.tokens.append(token)
 
             if self.logprobs > 0:
                 sorted_indices = mx.argpartition(-logprobs, kth=self.logprobs - 1)
@@ -549,9 +762,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(
@@ -561,23 +775,35 @@ 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)
+        if gen_response.finish_reason is not None:
+            finish_reason = gen_response.finish_reason
 
         logging.debug(f"Prompt: {gen_response.prompt_tps:.3f} tokens-per-sec")
         logging.debug(f"Generation: {gen_response.generation_tps:.3f} tokens-per-sec")
         logging.debug(f"Peak memory: {gen_response.peak_memory:.3f} GB")
 
         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"]:
-                response = self.completion_usage_response(len(prompt), len(tokens))
+                original_prompt_length = (
+                    len(self.prompt_cache.tokens) - len(tokens) + len(prompt)
+                )
+                response = self.completion_usage_response(
+                    original_prompt_length, len(tokens)
+                )
                 self.wfile.write(f"data: {json.dumps(response)}\n\n".encode())
                 self.wfile.flush()
             self.wfile.write("data: [DONE]\n\n".encode())
@@ -591,6 +817,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
@@ -640,8 +867,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"))
@@ -666,13 +894,25 @@ class APIHandler(BaseHTTPRequestHandler):
         """
         Respond to a GET request from a client.
         """
-        if self.path == "/v1/models":
+        if self.path.startswith("/v1/models"):
             self.handle_models_request()
+        elif self.path == "/health":
+            self.handle_health_check()
         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.
@@ -680,10 +920,27 @@ class APIHandler(BaseHTTPRequestHandler):
         self._set_completion_headers(200)
         self.end_headers()
 
+        files = ["config.json", "model.safetensors.index.json", "tokenizer_config.json"]
+
+        parts = self.path.split("/")
+        filter_repo_id = None
+        if len(parts) > 3:
+            filter_repo_id = "/".join(parts[3:])
+
+        def probably_mlx_lm(repo):
+            if repo.repo_type != "model":
+                return False
+            if "main" not in repo.refs:
+                return False
+            if filter_repo_id is not None and repo.repo_id != filter_repo_id:
+                return False
+            file_names = {f.file_path.name for f in repo.refs["main"].files}
+            return all(f in file_names for f in files)
+
         # Scan the cache directory for downloaded mlx models
         hf_cache_info = scan_cache_dir()
         downloaded_models = [
-            repo for repo in hf_cache_info.repos if "mlx" in repo.repo_id
+            repo for repo in hf_cache_info.repos if probably_mlx_lm(repo)
         ]
 
         # Create a list of available models
@@ -758,6 +1015,18 @@ def main():
         default=8080,
         help="Port for the HTTP server (default: 8080)",
     )
+    parser.add_argument(
+        "--draft-model",
+        type=str,
+        help="A model to be used for speculative decoding.",
+        default=None,
+    )
+    parser.add_argument(
+        "--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",
@@ -782,6 +1051,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(
@@ -792,4 +1097,8 @@ def main():
 
 
 if __name__ == "__main__":
+    print(
+        "Calling `python -m mlx_lm.server...` directly is deprecated."
+        " Use `mlx_lm.server...` or `python -m mlx_lm server ...` instead."
+    )
     main()

mlx_lm/tokenizer_utils.py

@@ -1,8 +1,9 @@
 import json
 from functools import partial
+from json import JSONDecodeError
 from typing import List
 
-from transformers import AutoTokenizer
+from transformers import AutoTokenizer, PreTrainedTokenizerFast
 
 
 class StreamingDetokenizer:
@@ -88,8 +89,9 @@ class NaiveStreamingDetokenizer(StreamingDetokenizer):
     def text(self):
         if self._current_tokens:
             self._current_text = self._tokenizer.decode(self._current_tokens)
-            if (
+            if self._current_text.endswith("\ufffd") or (
                 self._tokenizer.clean_up_tokenization_spaces
+                and len(self._current_text) > 0
                 and self._current_text[-1] == " "
             ):
                 self._current_text = self._current_text[:-1]
@@ -266,6 +268,26 @@ 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:
+            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
@@ -279,6 +301,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
@@ -301,6 +347,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
@@ -341,7 +416,9 @@ def _is_bpe_decoder(decoder):
     return isinstance(decoder, dict) and decoder.get("type", None) == "ByteLevel"
 
 
-def load_tokenizer(model_path, tokenizer_config_extra={}, eos_token_ids=None):
+def load_tokenizer(
+    model_path, tokenizer_config_extra={}, return_tokenizer=True, eos_token_ids=None
+):
     """Load a huggingface tokenizer and try to infer the type of streaming
     detokenizer to use.
 
@@ -353,7 +430,11 @@ def load_tokenizer(model_path, tokenizer_config_extra={}, eos_token_ids=None):
     tokenizer_file = model_path / "tokenizer.json"
     if tokenizer_file.exists():
         with open(tokenizer_file, "r", encoding="utf-8") as fid:
-            tokenizer_content = json.load(fid)
+            try:
+                tokenizer_content = json.load(fid)
+            except JSONDecodeError as e:
+                raise JSONDecodeError("Failed to parse tokenizer.json", e.doc, e.pos)
+
         if "decoder" in tokenizer_content:
             if _is_spm_decoder(tokenizer_content["decoder"]):
                 detokenizer_class = SPMStreamingDetokenizer
@@ -364,11 +445,15 @@ def load_tokenizer(model_path, tokenizer_config_extra={}, eos_token_ids=None):
 
     if isinstance(eos_token_ids, int):
         eos_token_ids = [eos_token_ids]
-    return TokenizerWrapper(
-        AutoTokenizer.from_pretrained(model_path, **tokenizer_config_extra),
-        detokenizer_class,
-        eos_token_ids=eos_token_ids,
-    )
+
+    if return_tokenizer:
+        return TokenizerWrapper(
+            AutoTokenizer.from_pretrained(model_path, **tokenizer_config_extra),
+            detokenizer_class,
+            eos_token_ids=eos_token_ids,
+        )
+    else:
+        return detokenizer_class
 
 
 def no_bos_or_eos(sequence: List, bos: int, eos: int) -> List:

mlx_lm/tuner/callbacks.py

@@ -0,0 +1,133 @@
+# Copyright © 2024 Apple Inc.
+
+import os
+
+try:
+    import wandb
+except ImportError:
+    wandb = None
+
+try:
+    import swanlab
+except ImportError:
+    swanlab = 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 wandb via: pip install wandb",
+            )
+        self.wrapped_callback = wrapped_callback
+        wandb.init(
+            project=project_name,
+            name=os.path.basename(log_dir),
+            dir=log_dir,
+            config=config,
+        )
+
+    def _convert_to_serializable(self, data: dict) -> dict:
+        return {k: v.tolist() if hasattr(v, "tolist") else v for k, v in data.items()}
+
+    def on_train_loss_report(self, train_info: dict):
+        wandb.log(
+            self._convert_to_serializable(train_info), step=train_info.get("iteration")
+        )
+        if self.wrapped_callback:
+            self.wrapped_callback.on_train_loss_report(train_info)
+
+    def on_val_loss_report(self, val_info: dict):
+        wandb.log(
+            self._convert_to_serializable(val_info), step=val_info.get("iteration")
+        )
+        if self.wrapped_callback:
+            self.wrapped_callback.on_val_loss_report(val_info)
+
+
+class SwanLabCallback(TrainingCallback):
+    def __init__(
+        self,
+        project_name: str,
+        log_dir: str,
+        config: dict,
+        wrapped_callback: TrainingCallback = None,
+    ):
+        if swanlab is None:
+            raise ImportError(
+                "swanlab is not installed. please install swanlab via: pip install swanlab",
+            )
+        self.wrapped_callback = wrapped_callback
+        swanlab.init(
+            project=project_name,
+            experiment_name=os.path.basename(log_dir),
+            logdir=os.path.join(log_dir, "swanlog"),
+            config=config,
+        )
+
+    def _convert_to_serializable(self, data: dict) -> dict:
+        return {k: v.tolist() if hasattr(v, "tolist") else v for k, v in data.items()}
+
+    def on_train_loss_report(self, train_info: dict):
+        swanlab.log(
+            self._convert_to_serializable(train_info), step=train_info.get("iteration")
+        )
+        if self.wrapped_callback:
+            self.wrapped_callback.on_train_loss_report(train_info)
+
+    def on_val_loss_report(self, val_info: dict):
+        swanlab.log(
+            self._convert_to_serializable(val_info), step=val_info.get("iteration")
+        )
+        if self.wrapped_callback:
+            self.wrapped_callback.on_val_loss_report(val_info)
+
+
+SUPPORT_CALLBACK = {
+    "wandb": WandBCallback,
+    "swanlab": SwanLabCallback,
+}
+
+
+def get_reporting_callbacks(
+    report_to: str = None,
+    project_name: str = None,
+    log_dir: str = None,
+    config: str = None,
+):
+    if report_to is None or report_to == "":
+        return None
+    report_to = [item.strip().lower() for item in report_to.split(",") if item.strip()]
+    training_callback = None
+    for callback in report_to:
+        try:
+            training_callback = SUPPORT_CALLBACK[callback](
+                project_name=project_name,
+                log_dir=log_dir,
+                config=config,
+                wrapped_callback=training_callback,
+            )
+        except KeyError as e:
+            raise ValueError(
+                f"{callback} callback doesn't exist "
+                f"choose from {', '.join(SUPPORT_CALLBACK.keys())}"
+            ) from e
+
+    return training_callback

mlx_lm/tuner/datasets.py

@@ -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
 
@@ -17,7 +19,7 @@ class TextDataset:
         tokenizer: PreTrainedTokenizer,
         text_key: str = "text",
     ):
-        self._data = [d for d in data]
+        self._data = data
         self.tokenizer = tokenizer
         self.text_key = text_key
 
@@ -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]
@@ -47,7 +49,7 @@ class ChatDataset:
         chat_key: str = "messages",
         mask_prompt: bool = False,
     ):
-        self._data = [d for d in data]
+        self._data = data
         self.chat_key = chat_key
         self.mask_prompt = mask_prompt
         self.tokenizer = tokenizer
@@ -58,13 +60,10 @@ class ChatDataset:
         tokens = self.tokenizer.apply_chat_template(messages, tools=tools)
         if self.mask_prompt:
             messages = messages[:-1]
-            offset = len(tokenizer.apply_chat_template(messages, tools=tools))
+            offset = len(self.tokenizer.apply_chat_template(messages, tools=tools))
             return (tokens, offset)
         else:
-            return tokens
-
-    def itemlen(idx: int):
-        return len(self._data[idx])
+            return (tokens, 0)
 
     def __getitem__(self, idx: int):
         return self._data[idx]
@@ -88,7 +87,7 @@ class CompletionsDataset:
         completion_key: str,
         mask_prompt: bool,
     ):
-        self._data = [d for d in data]
+        self._data = data
         self.prompt_key = prompt_key
         self.completion_key = completion_key
         self.mask_prompt = mask_prompt
@@ -109,7 +108,7 @@ class CompletionsDataset:
             )
             return (tokens, offset)
 
-        return tokens
+        return (tokens, 0)
 
     def __getitem__(self, idx: int):
         return self._data[idx]
@@ -124,12 +123,17 @@ class ConcatenatedDataset:
         self._len = sum(len(d) for d in self._data)
 
     def __getitem__(self, idx: int):
-        for data in self._data:
+        for data_idx, data in enumerate(self._data):
             j = idx - len(data)
             if j < 0:
                 break
             idx = j
-        return data[idx]
+        datum = data[idx]
+        datum["_dataset"] = data_idx
+        return datum
+
+    def process(self, d):
+        return self._data[d["_dataset"]].process(d)
 
     def __len__(self):
         return self._len
@@ -178,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."
         )
 
 

mlx_lm/tuner/lora.py

@@ -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,