fix for rnj-1 (#657)

Co-authored-by: Awni Hannun <awni@apple.com>

.circleci/config.yml

@@ -1,100 +0,0 @@
-version: 2.1
-
-orbs:
-  apple: ml-explore/pr-approval@0.1.0
- 
-jobs:
-  linux_build_and_test:
-    docker:
-      - image: cimg/python:3.9
-
-    steps:
-      - checkout
-      - run:
-          name: Run style checks
-          command: |
-            pip install pre-commit
-            pre-commit run --all
-            if ! git diff --quiet; then echo 'Style checks failed, please install pre-commit and run pre-commit run --all and push the change'; exit 1; fi
-
-  mlx_lm_build_and_test:
-    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 sentencepiece
-            pip install unittest-xml-reporting
-            pip install -e ".[test]"
-      - run:
-          name: Run Python tests
-          command: |
-            source env/bin/activate
-            python -m xmlrunner discover -v tests -o test-results/
-      - 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:
-      matches:
-        pattern: "^(?!pull/)[-\\w]+$"
-        value: << pipeline.git.branch >>
-    jobs:
-      - mlx_lm_build_and_test
-      - linux_build_and_test
-
-  build_pypi_release:
-    jobs:
-      - build_release:
-          filters:
-            tags:
-              only: /^v.*/
-            branches:
-              ignore: /.*/
-
-  prb:
-    when:
-      matches:
-        pattern: "^pull/\\d+(/head)?$"
-        value: << pipeline.git.branch >>
-    jobs:
-      - hold:
-          type: approval
-      - apple/authenticate:
-          context: pr-approval
-      - mlx_lm_build_and_test:
-          requires: [ hold ]
-      - linux_build_and_test:
-          requires: [ hold ]

.github/actions/setup-macos/action.yml

@@ -0,0 +1,16 @@
+name: 'Setup macOS Environment'
+description: 'Install dependencies for macOS'
+
+inputs:
+  python-version:
+    description: 'Python version to use'
+    required: false
+    default: '3.10'
+
+runs:
+  using: "composite"
+  steps:
+    - uses: conda-incubator/setup-miniconda@v3
+      with:
+        miniconda-version: "latest"
+        python-version: ${{ inputs.python-version }}

.github/workflows/pull_request.yml

@@ -0,0 +1,41 @@
+name: Build and Test
+
+on:
+  push:
+    branches: ["main"]
+  pull_request:
+
+permissions:
+  contents: read
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ github.ref != 'refs/head/main' }}
+
+jobs:
+  check_lint:
+    if: github.repository == 'ml-explore/mlx-lm'
+    runs-on: ubuntu-22.04
+    steps:
+      - uses: actions/checkout@v5
+      - uses: actions/setup-python@v6
+        with:
+          python-version: "3.10"
+      - uses: pre-commit/action@v3.0.1
+
+  mac_build_and_test:
+    if: github.repository == 'ml-explore/mlx-lm'
+    runs-on: [self-hosted, macos]
+    needs: check_lint
+    steps:
+      - uses: actions/checkout@v5
+      - uses: ./.github/actions/setup-macos
+      - name: Install test dependencies
+        shell: bash -l {0}
+        run: |
+          pip install unittest-xml-reporting
+          pip install -e ".[test]"
+      - name: Run tests
+        shell: bash -l {0}
+        run: |
+          python -m xmlrunner discover -v tests -o test-results/

.github/workflows/release.yml

@@ -0,0 +1,41 @@
+name: PyPI Release
+
+on:
+  push:
+    tags:
+      - 'v*'
+  workflow_dispatch:
+
+permissions:
+  contents: read
+
+jobs:
+
+  build_release:
+    if: github.repository == 'ml-explore/mlx-lm'
+    runs-on: ubuntu-22.04
+    permissions:
+      id-token: write
+    environment:
+      name: pypi
+      url: https://pypi.org/p/mlx-lm
+    steps:
+      - uses: actions/checkout@v5
+      - uses: actions/setup-python@v6
+        with:
+          python-version: "3.10"
+      - name: Build package
+        shell: sh
+        run: |
+          pip install build
+          python -m build
+      - name: Upload artifacts
+        uses: actions/upload-artifact@v5
+        with:
+          overwrite: true
+          name: mlx-lm
+          path: dist/*
+      - name: Publish package distributions to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+        with:
+          repository-url: https://upload.pypi.org/legacy/

ACKNOWLEDGMENTS.md

@@ -8,17 +8,22 @@ 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.
-- Gökdeniz Gülmez: Added support for the following architectures: OpenBMB's
-  `MiniCPM` and `MiniCPM3`, Kyutai's `Helium`, State-Space's `Mamba v1`, `Mamba v2`, Z.ai &
-   THUKEG's `GLM`, `GLM4`, Rednote `dots.llm1`, Baisu's `Ernie4.5 MoE`, inclusionAI's
-   `Bailing MoE e.g. Ling-family`, Klear team - Kuaishou Technology's `Klear`,
-   IBM's `Granite MoE`, Meituan's `LongCat`, Nvidia's `Nemotron H`, Swiss-AI's
-   `Apertus`, Nikity's `Lille130m`, Alibaba Qwen's `Qwen3Next`, and Allenai's `OLMoE`;
-   Helped add support for the following model architectures: Alibaba Qwen's `Qwen3 & Qwen3MoE)`;
-   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)`.
+- Gökdeniz Gülmez: Added support for the following architectures: 
+OpenBMB's `MiniCPM` and `MiniCPM3`, Kyutai's `Helium`, State-Space's `Mamba v1` and 
+`Mamba v2`, Z.ai & THUKEG's `GLM`, `GLM4`, Rednote `dots.llm1`, Baidu's `Ernie4.5 MoE`, 
+inclusionAI's `Bailing MoE e.g. Ling-family`, `Bailing MoE Linear e.g. Ling-Linear-family`, 
+Klear team - Kuaishou Technology's `Klear`, AI21 Lab's `Jamba` IBM's `Granite MoE`, 
+Meituan's `LongCat`, Nvidia's `Nemotron H`, Swiss-AI's `Apertus`, Nikity's `Lille130m`, 
+Alibaba Qwen's `Qwen3Next`, and Allenai's `OLMoE` and `Olmo 3`; 
+Helped add support for the following model architectures: 
+Alibaba Qwen's `Qwen3 & Qwen3MoE)`; 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`.
+  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)`, MinimaxAI's `MiniMax`, 
+  MoonshotAI's `Kimi-Linear`, LiquidAI's `LFM2` and `LFM2 MoE`, 
+  Google DeepMind's `Gemma 3`, TII's `Falcon H1` and InterLM's `InternLM 2.5`.
+- Ivan Fioravanti: Added support for the following architectures:
+ ServiceNow-AI's `Apriel 1.5`, Tencent's `Hunyuan Dense V1` and `Hunyuan MoE V1`.

README.md

@@ -236,45 +236,19 @@ for more usage details.
 
 ### Supported Models
 
-`mlx-lm` supports thousands of Hugging Face format LLMs. If the model you want to
-run is not supported, file an
-[issue](https://github.com/ml-explore/mlx-lm/issues/new) or better yet,
-submit a pull request.
+`mlx-lm` supports thousands of LLMs available on the Hugging Face Hub. If the
+model you want to run is not supported, file an
+[issue](https://github.com/ml-explore/mlx-lm/issues/new) or better yet, submit
+a pull request. Many supported models are available in various quantization
+formats in the [MLX Community](https://huggingface.co/mlx-community) Hugging
+Face organization.
 
-Here are a few examples of Hugging Face models that work with this example:
+For some models the tokenizer may require you to enable the `trust_remote_code`
+option. You can do this by passing `--trust-remote-code` in the command line.
+If you don't specify the flag explicitly, you will be prompted to trust remote
+code in the terminal when running the model. 
 
-- [mistralai/Mistral-7B-v0.1](https://huggingface.co/mistralai/Mistral-7B-v0.1)
-- [meta-llama/Llama-2-7b-hf](https://huggingface.co/meta-llama/Llama-2-7b-hf)
-- [deepseek-ai/deepseek-coder-6.7b-instruct](https://huggingface.co/deepseek-ai/deepseek-coder-6.7b-instruct)
-- [01-ai/Yi-6B-Chat](https://huggingface.co/01-ai/Yi-6B-Chat)
-- [microsoft/phi-2](https://huggingface.co/microsoft/phi-2)
-- [mistralai/Mixtral-8x7B-Instruct-v0.1](https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1)
-- [Qwen/Qwen-7B](https://huggingface.co/Qwen/Qwen-7B)
-- [pfnet/plamo-13b](https://huggingface.co/pfnet/plamo-13b)
-- [pfnet/plamo-13b-instruct](https://huggingface.co/pfnet/plamo-13b-instruct)
-- [stabilityai/stablelm-2-zephyr-1_6b](https://huggingface.co/stabilityai/stablelm-2-zephyr-1_6b)
-- [internlm/internlm2-7b](https://huggingface.co/internlm/internlm2-7b)
-- [tiiuae/falcon-mamba-7b-instruct](https://huggingface.co/tiiuae/falcon-mamba-7b-instruct)
-
-Most
-[Mistral](https://huggingface.co/models?library=transformers,safetensors&other=mistral&sort=trending),
-[Llama](https://huggingface.co/models?library=transformers,safetensors&other=llama&sort=trending),
-[Phi-2](https://huggingface.co/models?library=transformers,safetensors&other=phi&sort=trending),
-and
-[Mixtral](https://huggingface.co/models?library=transformers,safetensors&other=mixtral&sort=trending)
-style models should work out of the box.
-
-For some models (such as `Qwen` and `plamo`) the tokenizer requires you to
-enable the `trust_remote_code` option. You can do this by passing
-`--trust-remote-code` in the command line. If you don't specify the flag
-explicitly, you will be prompted to trust remote code in the terminal when
-running the model. 
-
-For `Qwen` models you must also specify the `eos_token`. You can do this by
-passing `--eos-token "<|endoftext|>"` in the command
-line. 
-
-These options can also be set in the Python API. For example:
+Tokenizer options can also be set in the Python API. For example:
 
 ```python
 model, tokenizer = load(

mlx_lm/BENCHMARKS.md

@@ -0,0 +1,63 @@
+# Benchmarks
+
+## Commands 
+
+The command for evaluating on MMLU Pro:
+
+```
+mlx_lm.evaluate --model model/repo --task mmlu_pro
+```
+ 
+The command for efficiency benchmarks:
+
+```
+mlx_lm.benchmark --model model/repo -p 2048 -g 128
+```
+
+To get the package versions run:
+
+```
+python -m mlx --version && python -m mlx_lm --version
+```
+
+## Models
+
+<details>
+
+ <summary> Qwen/Qwen3-4B-Instruct-2507 </summary>
+
+Precision | MMLU Pro | Prompt (2048) tok/sec | Generation (128) tok/sec | Memory GB | Repo
+--------- | -------- | ------------------- | ------------------------ | --------- | ----
+bf16      | 64.05    | 1780.63             | 52.47                    | 9.02    | Qwen/Qwen3-4B-Instruct-2507
+q8 | 63.85 | 1606.573| 86.907 | 5.254 | mlx-community/Qwen3-4B-Instruct-2507-8bit
+q6 | 63.53 | 1576.73 | 104.68 | 4.25 | mlx-community/Qwen3-4B-Instruct-2507-6bit
+q5 g32 | 63.16 | 1570.80 | 110.29 | 4.00 | mlx-community/Qwen3-4B-Instruct-2507-5bit-g32
+q5 | 62.38 | 1584.33 | 116.39 | 3.86 | mlx-community/Qwen3-4B-Instruct-2507-5bit
+q4 g32 | 61.46 | 1610.03 | 126.00 | 3.603 | mlx-community/Qwen3-4B-Instruct-2507-4bit-g32
+q4 | 60.72 | 1622.27 | 134.52 | 3.35 | mlx-community/Qwen3-4B-Instruct-2507-4bit
+
+- Performance benchmark on 64GB M4 Max
+- mlx 0.29.2.dev20251008+85a8824a8
+- mlx-lm 0.28.2
+- macOS 26.1
+ 
+</details>
+
+<details>
+<summary> Qwen/Qwen3-30B-A3B-Instruct-2507 </summary>
+
+Precision | MMLU Pro | Prompt (2048) tok/sec | Generation (128) tok/sec | Memory GB | Repo
+--------- | -------- | ------------------- | ------------------------ | --------- | ----
+bf16 | 72.62 | :skull: | :skull: | :skull: | Qwen/Qwen3-30B-A3B-Instruct-2507
+q8 | 72.46 | 1719.47 | 83.16 | 33.46 | mlx-community/Qwen3-30B-A3B-Instruct-2507-8bit 
+q6 | 72.41 | 1667.45 | 94.14 | 25.82 | mlx-community/Qwen3-30B-A3B-Instruct-2507-6bit
+q5 | 71.97 | 1664.24 | 101.00 |22.01 | mlx-community/Qwen3-30B-A3B-Instruct-2507-5bit
+q4 | 70.71 | 1753.90 | 113.33 |18.20 | mlx-community/Qwen3-30B-A3B-Instruct-2507-4bit
+
+ 
+- Performance benchmarks on 64GB M4 Max
+- mlx 0.29.2.dev20251008+85a8824a8
+- mlx-lm 0.28.2
+- macOS 26.1
+
+</details>

mlx_lm/LEARNED_QUANTS.md

@@ -129,7 +129,7 @@ mlx_lm.awq --help
 Use `mlx_lm.gptq` to run GPTQ on a given model. For example:
 
 ```bash
-mlx_lm.awq --model Qwen/Qwen3-0.6B
+mlx_lm.gptq --model Qwen/Qwen3-0.6B
 ```
 
 The script can take anywhere from a few minutes to several hours depending on

mlx_lm/LORA.md

@@ -371,7 +371,10 @@ of memory. Here are some tips to reduce memory use should you need to do so:
 
 2. Try using a smaller batch size with `--batch-size`. The default is `4` so
    setting this to `2` or `1` will reduce memory consumption. This may slow
-   things down a little, but will also reduce the memory use.
+   things down a little, but will also reduce the memory use. You can increase
+   the effective batch size without increasing the memory use by accumulating
+   gradients using `--grad-accumulation-steps <N>` which will accumulate the
+   gradient of `<N>` batches before updating the parameters.
 
 3. Reduce the number of layers to fine-tune with `--num-layers`. The default
    is `16`, so you can try `8` or `4`. This reduces the amount of memory

mlx_lm/__init__.py

@@ -9,3 +9,12 @@ os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1"
 from .convert import convert
 from .generate import batch_generate, generate, stream_generate
 from .utils import load
+
+__all__ = [
+    "__version__",
+    "convert",
+    "batch_generate",
+    "generate",
+    "stream_generate",
+    "load",
+]

mlx_lm/__main__.py

@@ -25,7 +25,12 @@ if __name__ == "__main__":
     if len(sys.argv) < 2:
         raise ValueError(f"CLI requires a subcommand in {subcommands}")
     subcommand = sys.argv.pop(1)
-    if subcommand not in subcommands:
+    if subcommand in subcommands:
+        submodule = importlib.import_module(f"mlx_lm.{subcommand}")
+        submodule.main()
+    elif subcommand == "--version":
+        from mlx_lm import __version__
+
+        print(__version__)
+    else:
         raise ValueError(f"CLI requires a subcommand in {subcommands}")
-    submodule = importlib.import_module(f"mlx_lm.{subcommand}")
-    submodule.main()

mlx_lm/_version.py

@@ -1,3 +1,3 @@
 # Copyright © 2023-2025 Apple Inc.
 
-__version__ = "0.28.0"
+__version__ = "0.29.0"

mlx_lm/benchmark.py

@@ -4,13 +4,9 @@ import argparse
 
 import mlx.core as mx
 
-from mlx_lm import batch_generate, stream_generate
+from mlx_lm import batch_generate, load, 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,
-)
+from mlx_lm.utils import pipeline_load
 
 
 def setup_arg_parser():
@@ -25,11 +21,6 @@ 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(
         "--prompt-tokens",
         "-p",
@@ -66,22 +57,30 @@ def main():
     args = parser.parse_args()
     mx.random.seed(0)
 
+    group = mx.distributed.init()
+    rank = group.rank()
+
+    def rprint(*args, **kwargs):
+        if rank == 0:
+            print(*args, **kwargs)
+
     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},
-    )
+    if group.size() > 1:
+        model, tokenizer, config = pipeline_load(args.model, return_config=True)
+    else:
+        model, tokenizer, config = load(
+            args.model, return_config=True, tokenizer_config={"trust_remote_code": True}
+        )
+
+    # Empty to avoid early stopping
+    tokenizer._eos_token_ids = {}
 
     prompt_tokens = args.prompt_tokens
     generation_tokens = args.generation_tokens
     batch_size = args.batch_size
-    prompts = mx.random.randint(
-        0, config["vocab_size"], (batch_size, prompt_tokens)
-    ).tolist()
+    vocab_size = config.get("vocab_size") or config["text_config"]["vocab_size"]
+    prompts = mx.random.randint(0, vocab_size, (batch_size, prompt_tokens)).tolist()
     prompt = prompts[0]
 
     def single_bench():
@@ -96,20 +95,23 @@ def main():
             model, tokenizer, prompts, max_tokens=generation_tokens
         ).stats
 
-    _bench = batch_bench
+    if batch_size == 1:
+        _bench = single_bench
+    else:
+        _bench = batch_bench
 
-    print("Running warmup..")
+    rprint("Running warmup..")
     _bench()
 
     report_keys = ["prompt_tps", "generation_tps", "peak_memory"]
-    print(f"Timing with {prompt_tokens=}, {generation_tokens=}, {batch_size=}.")
+    rprint(f"Timing with {prompt_tokens=}, {generation_tokens=}, {batch_size=}.")
     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))
+        rprint(f"Trial {i+1}:  " + ", ".join(results))
 
     def avg(k):
         vals = (getattr(response, k) for response in responses)
@@ -117,7 +119,7 @@ def main():
 
     results = [(k, avg(k)) for k in report_keys]
     results = [f"{k}={v:.3f}" for k, v in results]
-    print(f"Averages: " + ", ".join(results))
+    rprint(f"Averages: " + ", ".join(results))
 
 
 if __name__ == "__main__":

mlx_lm/convert.py

@@ -10,8 +10,7 @@ from mlx.utils import tree_map_with_path
 
 from .utils import (
     dequantize_model,
-    fetch_from_hub,
-    get_model_path,
+    load,
     quantize_model,
     save,
     upload_to_hub,
@@ -63,7 +62,9 @@ def mixed_quant_predicate_builder(
             or index >= 7 * num_layers // 8
             or (index - num_layers // 8) % 3 == 2
         )
-        if "v_proj" in path and use_more_bits:
+        if (
+            "v_proj" in path or "v_a_proj" in path or "v_b_proj" in path
+        ) and use_more_bits:
             return {"group_size": group_size, "bits": high_bits}
         if "down_proj" in path and use_more_bits:
             return {"group_size": group_size, "bits": high_bits}
@@ -107,9 +108,12 @@ def convert(
         )
 
     print("[INFO] Loading")
-    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
+    model, tokenizer, config = load(
+        hf_path,
+        revision=revision,
+        return_config=True,
+        tokenizer_config={"trust_remote_code": trust_remote_code},
+        lazy=True,
     )
 
     if isinstance(quant_predicate, str):
@@ -154,11 +158,10 @@ def convert(
 
     save(
         mlx_path,
-        model_path,
+        hf_path,
         model,
         tokenizer,
         config,
-        hf_repo=hf_path,
     )
 
     if upload_repo is not None:

mlx_lm/evaluate.py

@@ -12,7 +12,7 @@ import logging
 import os
 from importlib.metadata import version
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any, Callable, Optional
 
 import lm_eval
 import mlx.core as mx
@@ -25,7 +25,10 @@ from tqdm import tqdm
 
 from .generate import batch_generate
 from .models.cache import make_prompt_cache
-from .utils import common_prefix_len, load
+from .sample_utils import make_sampler
+from .utils import load
+
+DEFAULT_MAX_TOKENS = 8192
 
 
 def _rstrip_until(s, untils):
@@ -36,6 +39,13 @@ def _rstrip_until(s, untils):
     return s[: min(f)]
 
 
+def _lstrip(s, pattern):
+    """Truncate the prefix of the string after the first occurrence of pattern."""
+    if (idx := s.find(pattern)) != -1:
+        return s[idx + len(pattern) :]
+    return s
+
+
 def _pad_inputs(inputs):
     lengths = np.array([len(x) for x in inputs])
     maxlen = lengths.max()
@@ -68,9 +78,10 @@ class MLXLM(LM):
     def __init__(
         self,
         path_or_hf_repo: str,
-        max_tokens: int,
+        max_tokens: Optional[int] = None,
         use_chat_template: Optional[bool] = None,
         trust_remote_code: bool = False,
+        sampler: Optional[Callable[[mx.array], mx.array]] = None,
     ) -> None:
         super().__init__()
         tokenizer_config = {"trust_remote_code": True if trust_remote_code else None}
@@ -82,6 +93,7 @@ class MLXLM(LM):
         self.use_chat_template = use_chat_template
         if use_chat_template is None:
             self.use_chat_template = self.tokenizer.chat_template is not None
+        self._sampler = sampler
 
     def _process_prompt(self, prompt, step_size: int = 2048):
         prompt = mx.array(prompt)[None]
@@ -182,7 +194,8 @@ class MLXLM(LM):
             max_completed_l = max(len(s) for s in full_sequences)
 
             # compute truncation length
-            truncation = max(0, max_completed_l - self._max_tokens - 1)
+            max_tokens = self._max_tokens or DEFAULT_MAX_TOKENS
+            truncation = max(0, max_completed_l - max_tokens - 1)
             orig_prefix_l = len(prefix)
             prefix_l = max(len(prefix) - truncation, 0)
             prefix = prefix[len(prefix) - prefix_l :]
@@ -304,11 +317,16 @@ class MLXLM(LM):
             continuation: str
                 The generated continuation.
         """
+        group = mx.distributed.init()
+
+        # split data accross ranks
+        total_requests = len(requests)
+        requests = requests[group.rank() :: group.size()]
+
         logging.info("Generating continuation for %d sequences." % len(requests))
         contexts, options = zip(*[req.args for req in requests])
         # The second element of the tuple contains:
         # {'do_sample': False, 'until': ['\n\n'], 'temperature': 0}
-        completions = []
 
         # Tokenize all contexts
         contexts = [
@@ -319,7 +337,10 @@ class MLXLM(LM):
         ]
 
         # TODO consider multi-token, per-prompt stop conditions
-        max_tokens = [opt.get("max_gen_toks", self._max_tokens) for opt in options]
+        max_tokens = [
+            self._max_tokens or opt.get("max_gen_tokens", DEFAULT_MAX_TOKENS)
+            for opt in options
+        ]
 
         completions = batch_generate(
             model=self._model,
@@ -327,12 +348,46 @@ class MLXLM(LM):
             prompts=contexts,
             max_tokens=max_tokens,
             verbose=True,
+            sampler=self._sampler,
         ).texts
 
         for e, (text, opt) in enumerate(zip(completions, options)):
-            until = opt["until"]
-            if any(u in text for u in until):
-                completions[e] = _rstrip_until(text, until)
+            completions[e] = _rstrip_until(text, opt["until"])
+            if self.tokenizer.has_thinking:
+                completions[e] = _lstrip(text, self.tokenizer.think_end)
+
+        # Gather the completions
+        if group.size() > 1:
+            with mx.stream(mx.cpu):
+                pad_to = (total_requests + group.size() - 1) // group.size()
+                pad = pad_to - len(completions)
+                completions = [list(c.encode("utf-8")) for c in completions]
+                max_len = mx.array(max(len(c) for c in completions))
+                max_len = mx.distributed.all_max(max_len).item()
+                lengths = mx.array([len(c) for c in completions] + [0] * pad)
+                completions = mx.array(
+                    [c + [0] * (max_len - len(c)) for c in completions]
+                    + [[0] * max_len] * pad,
+                    mx.uint8,
+                )
+                completions = (
+                    mx.distributed.all_gather(completions[None])
+                    .swapaxes(0, 1)
+                    .flatten(0, 1)
+                    .tolist()
+                )
+                lengths = (
+                    mx.distributed.all_gather(lengths[None])
+                    .swapaxes(0, 1)
+                    .flatten(0, 1)
+                    .tolist()
+                )
+                completions = completions[:total_requests]
+                lengths = lengths[:total_requests]
+                completions = [
+                    bytearray(c[:l]).decode() for c, l in zip(completions, lengths)
+                ]
+
         return completions
 
 
@@ -350,8 +405,9 @@ def main():
     parser.add_argument(
         "--max-tokens",
         type=int,
-        help="Maximum number of tokens to generate.",
-        default=8912,
+        help="Maximum number of tokens to generate. When set, this value takes"
+        " precedence over task specific defaults.",
+        default=None,
     )
     parser.add_argument(
         "--limit",
@@ -393,7 +449,9 @@ def main():
         action="store_true",
         help="Enable trusting remote code for tokenizer",
     )
-
+    parser.add_argument("--temp", type=float, default=0.0, help="Sampling temperature")
+    parser.add_argument("--top-p", type=float, default=1.0, help="Sampling top-p")
+    parser.add_argument("--top-k", type=int, default=0, help="Sampling top-k")
     args = parser.parse_args()
 
     output_dir = Path(args.output_dir)
@@ -410,11 +468,17 @@ def main():
     if world.size() > 1 and world.rank() == 0:
         print(f"Evaluating with {world.size()} nodes")
 
+    sampler = make_sampler(
+        temp=args.temp,
+        top_p=args.top_p,
+        top_k=args.top_k,
+    )
     lm = MLXLM(
         args.model,
         max_tokens=args.max_tokens,
         use_chat_template=args.apply_chat_template,
         trust_remote_code=args.trust_remote_code,
+        sampler=sampler,
     )
     MLXLM.apply_chat_template = chat_template_fn(**args.chat_template_args)
 

mlx_lm/examples/batch_generate_response.py

@@ -26,7 +26,26 @@ prompts = [
 ]
 
 # Set `verbose=True` to see generation statistics
-result = batch_generate(model, tokenizer, prompts, verbose=False, max_tokens=128)
+result = batch_generate(
+    model, tokenizer, prompts, verbose=False, return_prompt_caches=True
+)
+print(result.texts[-1])
 
-# The returned result contains texts completions in the same order as prompts
-print(result.texts[0])
+prompts = [
+    "Could you summarize that?",
+    "And what about the sea?",
+    "Try again?",
+    "And Mt Olympus?",
+]
+prompts = [
+    tokenizer.apply_chat_template(
+        [{"role": "user", "content": p}],
+        add_generation_prompt=True,
+    )
+    for p in prompts
+]
+
+result = batch_generate(
+    model, tokenizer, prompts, verbose=False, prompt_caches=result.caches
+)
+print(result.texts[-1])

mlx_lm/examples/lora_config.yaml

@@ -1,5 +1,5 @@
 # The path to the local model directory or Hugging Face repo.
-model: "mlx-community/Llama-3.2-1B-Instruct"
+model: "mlx-community/Llama-3.2-1B-Instruct-bf16"
 
 # Whether or not to train (boolean)
 train: true
@@ -47,6 +47,9 @@ steps_per_report: 10
 # Number of training steps between validations.
 steps_per_eval: 200
 
+# Number of micro-steps to accumulate before each optimizer update.
+grad_accumulation_steps: 1
+
 # Load path to resume training with the given adapter weights.
 resume_adapter_file: null
 
@@ -90,4 +93,3 @@ lora_parameters:
 #  valid_split: "train[-100:]"
 #  prompt_feature: "text"
 #  completion_feature: "summary"
-

mlx_lm/examples/pipeline_generate.py

@@ -17,71 +17,11 @@ 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
-from huggingface_hub import snapshot_download
-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(
-        snapshot_download(
-            repo,
-            allow_patterns=allow_patterns,
-        )
-    )
-
-
-def shard_and_load(repo):
-    # Get model path with everything but weight safetensors
-    model_path = download(
-        args.model,
-        allow_patterns=["*.json", "*.py", "tokenizer.model", "*.tiktoken", "*.txt"],
-    )
-
-    # Lazy load and shard model to figure out
-    # which weights we need
-    model, config = load_model(model_path, lazy=True, strict=False)
-
-    group = mx.distributed.init()
-    rank = group.rank()
-    model.model.pipeline(group)
-
-    # Figure out which files we need for the local shard
-    with open(model_path / "model.safetensors.index.json", "r") as fid:
-        weight_index = json.load(fid)["weight_map"]
-
-    local_files = set()
-    for k, _ in tree_flatten(model.parameters()):
-        local_files.add(weight_index[k])
-
-    # Download weights for local shard
-    download(args.model, allow_patterns=local_files)
-
-    # Load and shard the model, and load the weights
-    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())
-
-    # Synchronize processes before generation to avoid timeout if downloading
-    # model for the first time.
-    mx.eval(mx.distributed.all_sum(mx.array(1.0), stream=mx.cpu))
-    return model, tokenizer
 
+from mlx_lm import stream_generate
+from mlx_lm.utils import pipeline_load
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="LLM pipelined inference example")
@@ -112,7 +52,7 @@ if __name__ == "__main__":
         if rank == 0:
             print(*args, **kwargs)
 
-    model, tokenizer = shard_and_load(args.model)
+    model, tokenizer = pipeline_load(args.model)
 
     messages = [{"role": "user", "content": args.prompt}]
     prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)

mlx_lm/fuse.py

@@ -4,10 +4,9 @@ from pathlib import Path
 from mlx.utils import tree_flatten, tree_unflatten
 
 from .gguf import convert_to_gguf
-from .tuner.utils import dequantize, load_adapters
 from .utils import (
-    fetch_from_hub,
-    get_model_path,
+    dequantize_model,
+    load,
     save,
     upload_to_hub,
 )
@@ -40,8 +39,8 @@ def parse_arguments() -> argparse.Namespace:
         default=None,
     )
     parser.add_argument(
-        "--de-quantize",
-        help="Generate a de-quantized model.",
+        "--dequantize",
+        help="Generate a dequantized model.",
         action="store_true",
     )
     parser.add_argument(
@@ -62,14 +61,12 @@ def main() -> None:
     print("Loading pretrained model")
     args = parse_arguments()
 
-    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)
+    model, tokenizer, config = load(
+        args.model, adapter_path=args.adapter_path, return_config=True
+    )
 
     fused_linears = [
-        (n, m.fuse(de_quantize=args.de_quantize))
+        (n, m.fuse(dequantize=args.dequantize))
         for n, m in model.named_modules()
         if hasattr(m, "fuse")
     ]
@@ -77,19 +74,18 @@ def main() -> None:
     if fused_linears:
         model.update_modules(tree_unflatten(fused_linears))
 
-    if args.de_quantize:
-        print("De-quantizing model")
+    if args.dequantize:
+        print("Dequantizing model")
         model = dequantize(model)
         config.pop("quantization", None)
 
     save_path = Path(args.save_path)
     save(
         save_path,
-        model_path,
+        args.model,
         model,
         tokenizer,
         config,
-        hf_repo=hf_path,
         donate_model=False,
     )
 
@@ -100,13 +96,9 @@ def main() -> None:
                 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))
+        convert_to_gguf(save_path, weights, config, str(save_path / args.gguf_path))
 
     if args.upload_repo is not 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)
 
 

mlx_lm/generate.py

@@ -7,6 +7,7 @@ import json
 import sys
 import time
 from dataclasses import dataclass
+from functools import partial
 from typing import (
     Any,
     Callable,
@@ -26,8 +27,11 @@ from .models import cache
 from .models.cache import (
     ArraysCache,
     BatchKVCache,
+    BatchRotatingKVCache,
+    CacheList,
     KVCache,
     QuantizedKVCache,
+    RotatingKVCache,
     load_prompt_cache,
 )
 from .sample_utils import make_sampler
@@ -284,16 +288,11 @@ class GenerationResponse:
 
 
 def maybe_quantize_kv_cache(prompt_cache, quantized_kv_start, kv_group_size, kv_bits):
-    if (
-        kv_bits is not None
-        and not isinstance(prompt_cache[0], cache.QuantizedKVCache)
-        and prompt_cache[0].offset > quantized_kv_start
-    ):
-        for i in range(len(prompt_cache)):
-            if isinstance(prompt_cache[i], cache.KVCache):
-                prompt_cache[i] = prompt_cache[i].to_quantized(
-                    group_size=kv_group_size, bits=kv_bits
-                )
+    if kv_bits is None:
+        return
+    for e, c in enumerate(prompt_cache):
+        if hasattr(c, "to_quantized") and c.offset >= quantized_kv_start:
+            prompt_cache[e] = c.to_quantized(group_size=kv_group_size, bits=kv_bits)
 
 
 def generate_step(
@@ -301,7 +300,7 @@ def generate_step(
     model: nn.Module,
     *,
     max_tokens: int = 256,
-    sampler: Optional[Callable[mx.array, mx.array]] = None,
+    sampler: Optional[Callable[[mx.array], mx.array]] = None,
     logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None,
     max_kv_size: Optional[int] = None,
     prompt_cache: Optional[Any] = None,
@@ -309,7 +308,7 @@ def generate_step(
     kv_bits: Optional[int] = None,
     kv_group_size: int = 64,
     quantized_kv_start: int = 0,
-    prompt_progress_callback: Optional[Callable[int, int]] = None,
+    prompt_progress_callback: Optional[Callable[[int, int], None]] = None,
     input_embeddings: Optional[mx.array] = None,
 ) -> Generator[Tuple[mx.array, mx.array], None, None]:
     """
@@ -335,7 +334,7 @@ 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_progress_callback (Callable[int, int]): A call-back which takes the
+        prompt_progress_callback (Callable[[int, int], None]): 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``.
@@ -420,7 +419,8 @@ def generate_step(
         prompt_processed_tokens = 0
         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)
+            remaining = (total_prompt_tokens - prompt_processed_tokens) - 1
+            n_to_process = min(prefill_step_size, remaining)
             _model_call(
                 input_tokens=prompt[:n_to_process][None],
                 input_embeddings=(
@@ -468,7 +468,7 @@ def speculative_generate_step(
     *,
     num_draft_tokens: int = 2,
     max_tokens: int = 256,
-    sampler: Optional[Callable[mx.array, mx.array]] = None,
+    sampler: Optional[Callable[[mx.array], mx.array]] = None,
     logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None,
     prompt_cache: Optional[Any] = None,
     prefill_step_size: int = 512,
@@ -487,7 +487,7 @@ def speculative_generate_step(
           speculative decoding. Default: ``2``.
         max_tokens (int): The maximum number of tokens. Use``-1`` for an infinite
           generator. Default: ``256``.
-        sampler (Callable[mx.array, mx.array], optional): A sampler for sampling a
+        sampler (Callable[[mx.array], mx.array], optional): A sampler for sampling a
           token from a vector of log probabilities. Default: ``None``.
         logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
           A list of functions that take tokens and logits and return the processed
@@ -787,6 +787,12 @@ def _left_pad_prompts(prompts, max_length=None):
     return mx.array([[0] * (max_length - len(p)) + p for p in prompts])
 
 
+def _right_pad_prompts(prompts, max_length=None):
+    if max_length is None:
+        max_length = max(len(p) for p in prompts)
+    return mx.array([p + [0] * (max_length - len(p)) for p in prompts])
+
+
 @dataclass
 class BatchStats:
     """
@@ -823,6 +829,7 @@ class BatchResponse:
 
     texts: List[str]
     stats: BatchStats
+    caches: Optional[List[List[Any]]]
 
 
 @dataclass
@@ -839,45 +846,71 @@ class Batch:
 
     def filter(self, keep_idx: List[int]):
         self.uids = [self.uids[k] for k in keep_idx]
+        self.logprobs = [self.logprobs[k] for k in keep_idx]
         self.max_tokens = [self.max_tokens[k] for k in keep_idx]
         self.num_tokens = [self.num_tokens[k] for k in keep_idx]
         keep_idx = mx.array(keep_idx, mx.int32)
         self.y = self.y[keep_idx]
-        self.logprobs = self.logprobs[keep_idx]
         for c in self.cache:
             c.filter(keep_idx)
 
     def extend(self, other):
         self.uids.extend(other.uids)
         self.y = mx.concatenate([self.y, other.y])
-        self.logprobs = mx.concatenate([self.logprobs, other.logprobs])
+        self.logprobs.extend(other.logprobs)
         self.num_tokens.extend(other.num_tokens)
         self.max_tokens.extend(other.max_tokens)
         for c, o in zip(self.cache, other.cache):
             c.extend(o)
 
+    def extract_cache(self, idx):
+        return [c.extract(idx) for c in self.cache]
+
 
 def _make_cache(model, left_padding):
     """
     Convert a list of regular caches into their corresponding
     batch-aware caches.
     """
+
+    def to_batch_cache(c):
+        if isinstance(c, KVCache):
+            return BatchKVCache(left_padding)
+        elif isinstance(c, ArraysCache):
+            c.left_padding = mx.array(left_padding)
+            return c
+        elif isinstance(c, RotatingKVCache):
+            if c.keep > 0:
+                raise ValueError("RotatingKVCache with keep tokens is not supported.")
+            return BatchRotatingKVCache(c.max_size, left_padding)
+        elif isinstance(c, CacheList):
+            return CacheList(*(to_batch_cache(sub_c) for sub_c in c.caches))
+        else:
+            raise ValueError(f"{type(c)} does not yet support batching")
+
     if hasattr(model, "make_cache"):
         cache = model.make_cache()
-        batch_cache = []
-        for c in cache:
-            if isinstance(c, KVCache):
-                batch_cache.append(BatchKVCache(left_padding))
-            elif isinstance(c, ArraysCache):
-                c.left_padding = mx.array(left_padding)
-                batch_cache.append(c)
-            else:
-                raise ValueError(f"{type(c)} does not yet support batching")
-        return batch_cache
+        return [to_batch_cache(c) for c in cache]
     else:
         return [BatchKVCache(left_padding) for _ in model.layers]
 
 
+def _merge_caches(caches):
+    batch_cache = []
+    for i in range(len(caches[0])):
+        cache = None
+        if isinstance(caches[0][i], KVCache):
+            cache = BatchKVCache.merge([c[i] for c in caches])
+        elif isinstance(caches[0][i], RotatingKVCache):
+            cache = BatchRotatingKVCache.merge([c[i] for c in caches])
+        else:
+            raise ValueError(
+                f"{type(caches[0][i])} does not yet support batching with history"
+            )
+        batch_cache.append(cache)
+    return batch_cache
+
+
 class BatchGenerator:
 
     @dataclass
@@ -886,16 +919,20 @@ class BatchGenerator:
         token: int
         logprobs: mx.array
         finish_reason: Optional[str]
+        prompt_cache: Callable[[], List[Any]]
 
     def __init__(
         self,
         model,
         max_tokens: int = 128,
         stop_tokens: Optional[set] = None,
-        sampler: Optional[Callable[mx.array, mx.array]] = None,
+        sampler: Optional[Callable[[mx.array], mx.array]] = None,
         completion_batch_size: int = 32,
         prefill_batch_size: int = 8,
         prefill_step_size: int = 2048,
+        prompt_progress_callback: Optional[
+            Callable[[List[Tuple[int, int, int]]], None]
+        ] = None,
     ):
         self.model = model
         self.unprocessed_prompts = []
@@ -905,44 +942,132 @@ class BatchGenerator:
         self.uid_count = 0
         self.prefill_step_size = prefill_step_size
         self.prefill_batch_size = prefill_batch_size
-        self.completion_batch_size = completion_batch_size
+        self.completion_batch_size = max(completion_batch_size, prefill_batch_size)
+        self.prompt_progress_callback = prompt_progress_callback or (lambda *_: None)
         self._stats = BatchStats()
 
         self.active_batch = None
 
-    def insert(self, prompts, max_tokens: Union[List[int], int, None] = None):
+        if mx.metal.is_available():
+            self._old_wired_limit = mx.set_wired_limit(
+                mx.metal.device_info()["max_recommended_working_set_size"]
+            )
+        else:
+            self._old_wired_limit = None
+
+    def close(self):
+        if self._old_wired_limit is not None:
+            mx.synchronize(generation_stream)
+            mx.set_wired_limit(self._old_wired_limit)
+            self._old_wired_limit = None
+
+    def __del__(self):
+        self.close()
+
+    def insert(
+        self, prompts, max_tokens: Union[List[int], int, None] = None, caches=None
+    ):
         uids = []
 
         if max_tokens is None or isinstance(max_tokens, int):
             max_tokens = [max_tokens or self.max_tokens] * len(prompts)
 
-        for p, m in zip(prompts, max_tokens):
-            self.unprocessed_prompts.append((self.uid_count, p, m))
+        if caches is None:
+            caches = [None] * len(prompts)
+        for i in range(len(prompts)):
+            if caches[i] is None:
+                caches[i] = cache.make_prompt_cache(self.model)
+
+        for p, m, c in zip(prompts, max_tokens, caches):
+            self.unprocessed_prompts.append((self.uid_count, p, m, c))
             uids.append(self.uid_count)
             self.uid_count += 1
         # Sort in ascending order of length
         self.unprocessed_prompts = sorted(
-            self.unprocessed_prompts, key=lambda x: len(x[1])
+            self.unprocessed_prompts, key=lambda x: len(x[1]) + cache.cache_length(x[3])
         )
         return uids
 
+    def remove(self, uids: List[int]):
+        uids = set(uids)
+        if self.active_batch is not None:
+            batch = self.active_batch
+            keep_idx = [e for e, uid in enumerate(batch.uids) if uid not in uids]
+            if len(keep_idx) > 0:
+                batch.filter(keep_idx)
+            else:
+                self.active_batch = None
+
+        for i in reversed(range(len(self.unprocessed_prompts))):
+            if self.unprocessed_prompts[i][0] in uids:
+                self.unprocessed_prompts.pop(i)
+
     def _process_prompts(self, prompts):
-        uids, inputs, max_tokens = zip(*prompts)
+        uids, inputs, max_tokens, caches = zip(*prompts)
+
+        cache_lengths = [cache.cache_length(c) for c in caches]
+        max_cache_length = max(cache_lengths)
         lengths = [len(p) for p in inputs]
         max_length = max(lengths)
-        batch_size = self.prefill_batch_size
+        padding = [max_length - l for l in lengths]
+
         self._stats.prompt_tokens += sum(lengths)
-        left_padding = [max_length - l for l in lengths]
-        inputs = _left_pad_prompts(inputs, max_length=max_length)
 
-        prompt_cache = _make_cache(self.model, left_padding)
+        processed_tokens = 0
 
-        while inputs.shape[1] > 1:
-            n_to_process = min(self.prefill_step_size, inputs.shape[1] - 1)
-            self.model(inputs[:, :n_to_process], cache=prompt_cache)
+        # New prompts so
+        #   1. Left-pad the inputs
+        #   2. Process
+        if max_cache_length == 0:
+            inputs = _left_pad_prompts(inputs, max_length=max_length)
+            prompt_cache = _make_cache(self.model, padding)
+
+            while inputs.shape[1] > 1:
+                n_to_process = min(self.prefill_step_size, inputs.shape[1] - 1)
+                self.model(inputs[:, :n_to_process], cache=prompt_cache)
+                mx.eval([c.state for c in prompt_cache])
+                inputs = inputs[:, n_to_process:]
+                processed_tokens += n_to_process
+                self.prompt_progress_callback(
+                    [
+                        (uid, processed_tokens, length)
+                        for uid, length in zip(uids, lengths)
+                    ]
+                )
+                mx.clear_cache()
+
+        # Further prompt processing so we need to
+        #   1. Merge the KV caches and prepare for right padded prompts
+        #   2. Right pad the inputs
+        #   2. Process
+        #   3. Finalize the KV caches so they are left padded again
+        else:
+            last_inputs = mx.array([p[-1:] for p in inputs])
+            inputs = _right_pad_prompts(inputs, max_length=max_length)
+            prompt_cache = _merge_caches(caches)
+
+            for c in prompt_cache:
+                c.prepare(lengths=lengths, right_padding=padding)
+
+            while inputs.shape[1] > 1:
+                n_to_process = min(self.prefill_step_size, inputs.shape[1] - 1)
+                self.model(inputs[:, :n_to_process], cache=prompt_cache)
+                mx.eval([c.state for c in prompt_cache])
+                inputs = inputs[:, n_to_process:]
+                processed_tokens += n_to_process
+                self.prompt_progress_callback(
+                    [
+                        (uid, processed_tokens, length)
+                        for uid, length in zip(uids, lengths)
+                    ]
+                )
+                mx.clear_cache()
+
+            for c in prompt_cache:
+                c.finalize()
             mx.eval([c.state for c in prompt_cache])
-            inputs = inputs[:, n_to_process:]
             mx.clear_cache()
+            inputs = last_inputs
 
         y, logprobs = self._step(inputs, prompt_cache)
         mx.async_eval(y, logprobs)
@@ -955,7 +1080,7 @@ class BatchGenerator:
         logits = logits[:, -1, :]
         logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
         sampled = self.sampler(logprobs)
-        return sampled, logprobs
+        return sampled, list(logprobs)
 
     def stats(self):
         self._stats.prompt_tps = self._stats.prompt_tokens / self._stats.prompt_time
@@ -1020,6 +1145,7 @@ class BatchGenerator:
         for e, (t, uid, num_tok, max_tok) in enumerate(
             zip(y, batch.uids, batch.num_tokens, batch.max_tokens)
         ):
+            cache = None
             num_tok += 1
             batch.num_tokens[e] = num_tok
             if t in self.stop_tokens:
@@ -1031,7 +1157,9 @@ class BatchGenerator:
             else:
                 finish_reason = None
                 keep_idx.append(e)
-            responses.append(self.Response(uid, t, logprobs[e], finish_reason))
+            if finish_reason is not None:
+                cache = batch.extract_cache(e)
+            responses.append(self.Response(uid, t, logprobs[e], finish_reason, cache))
 
         # Remove any finished completions
         if len(end_idx):
@@ -1052,8 +1180,10 @@ def batch_generate(
     model,
     tokenizer,
     prompts: List[int],
+    prompt_caches: Optional[List[List[Any]]] = None,
     max_tokens: Union[int, List[int]] = 128,
     verbose: bool = False,
+    return_prompt_caches: bool = False,
     **kwargs,
 ) -> BatchResponse:
     """
@@ -1063,10 +1193,15 @@ def batch_generate(
        model (nn.Module): The language model.
        tokenizer (PreTrainedTokenizer): The tokenizer.
        prompt (List[List[int]]): The input prompts.
+       prompt_caches (List[List[Any]], optional): Pre-computed prompt-caches
+          for each input prompt. Note, unlike ``generate_step``, the caches
+          won't be updated in-place.
        verbose (bool): If ``True``, print tokens and timing information.
           Default: ``False``.
        max_tokens (Union[int, List[int]): Maximum number of output tokens. This
           can be per prompt if a list is provided.
+       return_prompt_caches (bool): Return the prompt caches in the batch
+          responses. Default: ``False``.
        kwargs: The remaining options get passed to :obj:`BatchGenerator`.
           See :obj:`BatchGenerator` for more details.
     """
@@ -1077,25 +1212,30 @@ def batch_generate(
     if verbose:
         print(f"[batch_generate] Finished processing 0/{num_samples} ...", end="\r")
 
-    with wired_limit(model, [generation_stream]):
-        uids = gen.insert(prompts, max_tokens)
-        results = {uid: [] for uid in uids}
-        while responses := gen.next():
-            for r in responses:
-                if verbose and r.finish_reason != None:
+    uids = gen.insert(prompts, max_tokens, caches=prompt_caches)
+    results = {uid: [] for uid in uids}
+    prompt_caches = {}
+    while responses := gen.next():
+        for r in responses:
+            if r.finish_reason is not None:
+                if return_prompt_caches:
+                    prompt_caches[r.uid] = r.prompt_cache
+                if verbose:
                     fin += 1
                     print(
                         f"[batch_generate] Finished processing {fin}/{num_samples} ...",
                         end="\r",
                     )
-                if r.finish_reason != "stop":
-                    results[r.uid].append(r.token)
+            if r.finish_reason != "stop":
+                results[r.uid].append(r.token)
+    gen.close()
     if verbose:
         print(f"[batch_generate] Finished processing {fin}/{num_samples}")
 
     # Return results in correct order
     texts = [tokenizer.decode(results[uid]) for uid in uids]
     stats = gen.stats()
+    caches = [prompt_caches[uid] for uid in uids] if return_prompt_caches else None
     if verbose:
         print(
             f"[batch_generate] Prompt: {stats.prompt_tokens} tokens, {stats.prompt_tps:.3f} tokens-per-sec"
@@ -1105,7 +1245,7 @@ def batch_generate(
             f"{stats.generation_tps:.3f} tokens-per-sec"
         )
         print(f"[batch_generate] Peak memory: {stats.peak_memory:.3f} GB")
-    return BatchResponse(texts, stats)
+    return BatchResponse(texts, stats, caches)
 
 
 def main():

mlx_lm/lora.py

@@ -40,7 +40,7 @@ yaml_loader.add_implicit_resolver(
 )
 
 CONFIG_DEFAULTS = {
-    "model": "mlx_model",
+    "model": "Qwen/Qwen3-0.6b",
     "train": False,
     "fine_tune_type": "lora",
     "optimizer": "adam",
@@ -51,7 +51,7 @@ CONFIG_DEFAULTS = {
         "sgd": {},
         "adafactor": {},
     },
-    "data": "data/",
+    "data": "mlx-community/WikiSQL",
     "seed": 0,
     "num_layers": 16,
     "batch_size": 4,
@@ -68,10 +68,10 @@ CONFIG_DEFAULTS = {
     "max_seq_length": 2048,
     "config": None,
     "grad_checkpoint": False,
+    "grad_accumulation_steps": 1,
     "lr_schedule": None,
     "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,
 }
@@ -142,6 +142,11 @@ def build_parser():
         type=int,
         help="Number of training steps between validations.",
     )
+    parser.add_argument(
+        "--grad-accumulation-steps",
+        type=int,
+        help="Number of steps to accumulate before each optimizer update.",
+    )
     parser.add_argument(
         "--resume-adapter-file",
         type=str,
@@ -185,15 +190,6 @@ 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,
@@ -265,6 +261,7 @@ def train_model(
         adapter_file=adapter_file,
         max_seq_length=args.max_seq_length,
         grad_checkpoint=args.grad_checkpoint,
+        grad_accumulation_steps=args.grad_accumulation_steps,
     )
 
     # Initialize the selected optimizer
@@ -314,14 +311,6 @@ def evaluate_model(args, model: nn.Module, test_set):
 
 def run(args, training_callback: TrainingCallback = None):
     np.random.seed(args.seed)
-    if args.wandb is not None:
-        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,

mlx_lm/models/afm7.py

@@ -50,7 +50,7 @@ class FusedLoRALinear(nn.Module):
         ]
         self.lora_b = [mx.zeros((r, od)) for od in output_dims]
 
-    def fuse(self, de_quantize: bool = False):
+    def fuse(self, dequantize: bool = False):
         linear = self.linear
         weight = linear.weight
         is_quantized = isinstance(linear, FusedQuantizedLinear)
@@ -79,7 +79,7 @@ class FusedLoRALinear(nn.Module):
         delta = mx.concatenate(deltas, axis=0)
         fused_linear.weight = weight + delta
 
-        if is_quantized and not de_quantize:
+        if is_quantized and not dequantize:
             fused_linear = fused_linear.to_quantized(linear.group_size, linear.bits)
 
         return fused_linear

mlx_lm/models/afmoe.py

@@ -0,0 +1,404 @@
+# Copyright © 2024 Apple Inc.
+
+import math
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import KVCache, RotatingKVCache
+from .rope_utils import initialize_rope
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    layer_types: List[str]
+    vocab_size: int = 200192
+    hidden_size: int = 2048
+    intermediate_size: int = 6144
+    moe_intermediate_size: int = 1024
+    num_hidden_layers: int = 32
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 4
+    head_dim: int = 64
+    max_position_embeddings: int = 131072
+    rms_norm_eps: float = 1e-5
+    rope_theta: float = 10000
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = False
+    # MoE config
+    num_experts: int = 128
+    num_experts_per_tok: int = 8
+    num_shared_experts: int = 1
+    num_dense_layers: int = 2
+    route_norm: bool = True
+    route_scale: float = 2.826
+    score_func: str = "sigmoid"
+    n_group: int = 1
+    topk_group: int = 1
+    sliding_window: int = 2048
+    mup_enabled: bool = True
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs, is_local_attention: bool = False):
+        super().__init__()
+
+        self.hidden_size = args.hidden_size
+        self.n_heads = args.num_attention_heads
+        self.n_kv_heads = args.num_key_value_heads
+        self.head_dim = args.head_dim
+        self.is_local_attention = is_local_attention
+
+        self.scale = self.head_dim**-0.5
+
+        self.q_proj = nn.Linear(
+            self.hidden_size, self.n_heads * self.head_dim, bias=False
+        )
+        self.k_proj = nn.Linear(
+            self.hidden_size, self.n_kv_heads * self.head_dim, bias=False
+        )
+        self.v_proj = nn.Linear(
+            self.hidden_size, self.n_kv_heads * self.head_dim, bias=False
+        )
+        self.o_proj = nn.Linear(
+            self.n_heads * self.head_dim, self.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.gate_proj = nn.Linear(
+            self.hidden_size, self.n_heads * self.head_dim, bias=False
+        )
+
+        if is_local_attention:
+            self.rope = initialize_rope(
+                self.head_dim,
+                args.rope_theta,
+                False,  # traditional
+                args.rope_scaling,
+                args.max_position_embeddings,
+            )
+        else:
+            self.rope = None
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries = self.q_proj(x)
+        keys = self.k_proj(x)
+        values = 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
+        )
+
+        queries = self.q_norm(queries)
+        keys = self.k_norm(keys)
+
+        if self.is_local_attention and self.rope is not None:
+            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 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)
+
+        gate = mx.sigmoid(self.gate_proj(x))
+        output = output * gate
+
+        return self.o_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, args: ModelArgs, intermediate_size: Optional[int] = None):
+        super().__init__()
+
+        dim = args.hidden_size
+        hidden_dim = (
+            intermediate_size
+            if intermediate_size is not None
+            else 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 MoERouter(nn.Module):
+    """Router module that wraps the gate for proper weight naming."""
+
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.gate = nn.Linear(args.hidden_size, args.num_experts, bias=False)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.gate(x)
+
+
+class AfmoeMoE(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.num_experts = args.num_experts
+        self.num_experts_per_tok = args.num_experts_per_tok
+        self.route_norm = args.route_norm
+        self.route_scale = args.route_scale
+        self.score_func = args.score_func
+        self.n_group = args.n_group
+        self.topk_group = args.topk_group
+
+        self.router = MoERouter(args)
+
+        self.expert_bias = mx.zeros((args.num_experts,))
+
+        self.experts = SwitchGLU(
+            args.hidden_size,
+            args.moe_intermediate_size,
+            args.num_experts,
+        )
+
+        if args.num_shared_experts > 0:
+            shared_intermediate_size = (
+                args.moe_intermediate_size * args.num_shared_experts
+            )
+            self.shared_experts = MLP(args, intermediate_size=shared_intermediate_size)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        gates = self.router(x)
+
+        if self.score_func == "sigmoid":
+            scores = mx.sigmoid(gates.astype(mx.float32))
+        else:
+            scores = mx.softmax(gates.astype(mx.float32), axis=-1)
+
+        # Add expert bias for selection
+        selection_scores = scores + self.expert_bias
+
+        # Group-based expert selection if n_group > 1
+        if self.n_group > 1:
+            selection_scores = mx.unflatten(
+                selection_scores, axis=-1, shape=(self.n_group, -1)
+            )
+            group_scores = mx.topk(selection_scores, 2, axis=-1).sum(
+                axis=-1, keepdims=True
+            )
+            k = self.n_group - self.topk_group
+            group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :]
+            selection_scores = mx.put_along_axis(
+                selection_scores, mx.stop_gradient(group_idx), mx.array(0.0), axis=-2
+            )
+            selection_scores = mx.flatten(selection_scores, -2, -1)
+
+        # Select top-k experts
+        k = self.num_experts_per_tok
+        inds = mx.argpartition(-selection_scores, kth=k - 1, axis=-1)[..., :k]
+
+        selected_scores = mx.take_along_axis(scores, inds, axis=-1)
+
+        if self.route_norm and self.num_experts_per_tok > 1:
+            denominator = selected_scores.sum(axis=-1, keepdims=True)
+            selected_scores = selected_scores / denominator
+
+        selected_scores = selected_scores * self.route_scale
+
+        y = self.experts(x, inds)
+        y = (y * selected_scores[..., None]).sum(axis=-2).astype(y.dtype)
+
+        if self.args.num_shared_experts > 0:
+            y = y + self.shared_experts(x)
+
+        return y
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int, use_sliding: bool = False):
+        super().__init__()
+        self.hidden_size = args.hidden_size
+        self.use_sliding = use_sliding
+        self.layer_idx = layer_idx
+
+        self.self_attn = Attention(args, is_local_attention=use_sliding)
+
+        if layer_idx < args.num_dense_layers:
+            self.mlp = MLP(args)
+        else:
+            self.mlp = AfmoeMoE(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.pre_mlp_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:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        r = self.post_attention_layernorm(r)
+        h = x + r
+
+        r = self.mlp(self.pre_mlp_layernorm(h))
+        r = self.post_mlp_layernorm(r)
+        return h + r
+
+
+class AfmoeModel(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.layer_types = args.layer_types
+        self.sliding_window = args.sliding_window
+        self.mup_enabled = args.mup_enabled
+        self.hidden_size = args.hidden_size
+
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            DecoderLayer(
+                args=args, layer_idx=idx, use_sliding=layer_type == "sliding_attention"
+            )
+            for idx, layer_type in enumerate(self.layer_types)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+        self.fa_idx = self.layer_types.index("full_attention")
+        self.swa_idx = None
+        for idx, layer in enumerate(self.layers):
+            if layer.use_sliding:
+                self.swa_idx = idx
+                break
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+    ):
+        h = self.embed_tokens(inputs)
+
+        if self.mup_enabled:
+            h = h * math.sqrt(self.hidden_size)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        fa_mask = create_attention_mask(h, cache[self.fa_idx])
+        swa_mask = None
+        if self.swa_idx is not None:
+            swa_mask = create_attention_mask(
+                h, cache[self.swa_idx], window_size=self.sliding_window
+            )
+
+        for layer, c in zip(self.layers, cache):
+            mask = swa_mask if layer.use_sliding else fa_mask
+            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 = AfmoeModel(args)
+
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+    ):
+        out = self.model(inputs, 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 "rotary_emb.inv_freq" not in k}
+
+        if self.args.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
+
+        # Stack experts weights for SwitchGLU
+        for l in range(self.args.num_hidden_layers):
+            if l < self.args.num_dense_layers:
+                continue
+            prefix = f"model.layers.{l}"
+            for n in ["up_proj", "down_proj", "gate_proj"]:
+                for k in ["weight", "scales", "biases"]:
+                    if f"{prefix}.mlp.experts.0.{n}.{k}" in weights:
+                        to_join = [
+                            weights.pop(f"{prefix}.mlp.experts.{e}.{n}.{k}")
+                            for e in range(self.args.num_experts)
+                        ]
+                        weights[f"{prefix}.mlp.experts.{n}.{k}"] = mx.stack(to_join)
+
+        return weights
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def make_cache(self):
+        return [
+            (
+                RotatingKVCache(max_size=self.model.sliding_window)
+                if layer.use_sliding
+                else KVCache()
+            )
+            for layer in self.layers
+        ]
+
+    @property
+    def cast_predicate(self):
+        def predicate(k):
+            return "expert_bias" not in k
+
+        return predicate
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if "router.gate" in path:
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate

mlx_lm/models/bailing_moe.py

@@ -1,6 +1,7 @@
 # Copyright © 2025 Apple Inc.
 
 from dataclasses import dataclass
+from functools import partial
 from typing import Any, Dict, Optional, Union
 
 import mlx.core as mx
@@ -37,6 +38,7 @@ class ModelArgs(BaseModelArgs):
     use_qk_norm: bool = False
     tie_word_embeddings: bool = False
     partial_rotary_factor: float = 1.0
+    rotary_dim: Optional[int] = None
     moe_router_enable_expert_bias: bool = False
     moe_router_enable_routed_scaling: bool = True
     routed_scaling_factor: float = 1.0
@@ -47,6 +49,18 @@ class ModelArgs(BaseModelArgs):
     moe_router_enable_shared_expert: bool = True
 
 
+@partial(mx.compile, shapeless=True)
+def swiglu(gate, up):
+    return nn.silu(gate) * up
+
+
+@partial(mx.compile, shapeless=True)
+def aggregate_expert_outputs(expert_outputs, scores):
+    return (
+        (expert_outputs * scores[..., None]).sum(axis=-2).astype(expert_outputs.dtype)
+    )
+
+
 class BailingMoeMLP(nn.Module):
     def __init__(self, args: ModelArgs, intermediate_size: Optional[int] = None):
         super().__init__()
@@ -67,7 +81,7 @@ class BailingMoeMLP(nn.Module):
         )
 
     def __call__(self, x) -> mx.array:
-        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class BailingMoeAttention(nn.Module):
@@ -94,8 +108,10 @@ class BailingMoeAttention(nn.Module):
             self.key_layernorm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
             self.query_layernorm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
 
+        if (rope_dim := args.rotary_dim) is None:
+            rope_dim = int(self.head_dim * args.partial_rotary_factor)
         self.rope = initialize_rope(
-            int(self.head_dim * args.partial_rotary_factor),
+            rope_dim,
             args.rope_theta,
             traditional=False,
             scaling_config=args.rope_scaling,
@@ -146,6 +162,7 @@ class BailingMoeAttention(nn.Module):
         return self.dense(output)
 
 
+@mx.compile
 def group_expert_select(
     gates,
     e_score_correction_bias,
@@ -171,15 +188,15 @@ def group_expert_select(
         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.stop_gradient(group_idx), mx.array(0.0, scores.dtype), axis=-2
         )
         scores = mx.flatten(scores, -2, -1)
 
     k = top_k
-    inds = mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k]
+    inds = mx.argpartition(scores, kth=-k, 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)
+        denominator = scores.sum(axis=-1, keepdims=True) + 1e-20
         scores = scores / denominator
     scores = scores * routed_scaling_factor
 
@@ -245,7 +262,7 @@ class BailingMoeSparseMoeBlock(nn.Module):
     def __call__(self, x):
         topk_idx, topk_weight = self.gate(x)
         out = self.switch_mlp(x, topk_idx)
-        out = (out * topk_weight[..., None]).sum(axis=-2)
+        out = aggregate_expert_outputs(out, topk_weight)
         if self.shared_experts is not None:
             out = out + self.shared_experts(x)
         return out

mlx_lm/models/bailing_moe_linear.py

@@ -0,0 +1,594 @@
+# Copyright © 2025 Apple Inc.
+
+import math
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Tuple, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import (
+    BaseModelArgs,
+    create_attention_mask,
+    create_ssm_mask,
+    scaled_dot_product_attention,
+)
+from .cache import ArraysCache, KVCache
+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
+    layer_group_size: int
+    group_norm_size: int
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    rope_traditional: bool = False
+    use_bias: bool = False
+    use_qkv_bias: bool = False
+    norm_head: bool = False
+    norm_softmax: bool = False
+    use_qk_norm: bool = False
+    tie_word_embeddings: bool = False
+    partial_rotary_factor: float = 1.0
+    moe_router_enable_expert_bias: bool = False
+    moe_router_enable_routed_scaling: bool = True
+    routed_scaling_factor: float = 1.0
+    score_function: str = "softmax"
+    n_group: int = 1
+    topk_group: int = 4
+    use_rmsnorm: bool = True
+    moe_shared_expert_intermediate_size: Optional[int] = None
+    moe_router_enable_shared_expert: bool = True
+    head_dim: Optional[int] = None
+
+
+def recurrent_gla(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    scale: float,
+    h: Optional[mx.array] = None,
+) -> mx.array:
+    """
+    Recurrence per (b, h):
+        h_t = h_{t-1} * exp(g_t)
+        h_t = h_t + k_t^T @ v_t
+        y_t = (q_t @ h_t) * scale
+    Returns y with shape [B, H, T, Dv].
+    """
+    B, Hq, L, K = q.shape
+    Hv = k.shape[1]
+    V = v.shape[-1]
+
+    outputs = []
+    exp_g = mx.exp(g)[:, None, None].astype(q.dtype)
+    q = q * scale
+    for t in range(L):
+        q_t = q[:, :, t : t + 1]
+        k_t = k[:, :, t : t + 1]
+        v_t = v[:, :, t : t + 1]
+        h_up = k_t.transpose(0, 1, 3, 2) @ v_t
+        if h is not None:
+            h = h * exp_g + h_up
+        else:
+            h = h_up
+        o_t = q_t @ h
+        outputs.append(o_t)
+
+    return mx.concatenate(outputs, axis=2), h
+
+
+class GroupRMSNorm(nn.Module):
+    def __init__(self, dims: int, eps: float = 1e-5, groups: int = 1):
+        super().__init__()
+        self.weight = mx.ones((dims,))
+        self.groups = groups
+        self.eps = eps
+
+    def __call__(self, x: mx.array) -> mx.array:
+        shape = x.shape
+        x = mx.unflatten(x, axis=-1, shape=(self.groups, -1))
+        x = mx.fast.rms_norm(x, weight=None, eps=self.eps)
+        return self.weight * mx.flatten(x, -2)
+
+
+class MLP(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 Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.use_qk_norm = args.use_qk_norm
+        self.num_attention_heads = args.num_attention_heads
+        self.num_key_value_heads = args.num_key_value_heads
+        self.head_dim = args.head_dim or 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,
+        )
+
+        if args.use_qk_norm:
+            self.key_layernorm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
+            self.query_layernorm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
+
+        self.rope = initialize_rope(
+            int(self.head_dim * args.partial_rotary_factor),
+            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
+
+        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 self.use_qk_norm:
+            queries = self.query_layernorm(queries)
+            keys = self.key_layernorm(keys)
+
+        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 LinearAttention(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.use_qk_norm = args.use_qk_norm
+        self.num_hidden_layers = args.num_hidden_layers
+        self.num_attention_heads = args.num_attention_heads
+        self.num_key_value_heads = args.num_attention_heads
+        self.head_dim = args.hidden_size // self.num_attention_heads
+        self.scale = self.head_dim**-0.5
+        self.num_key_value_groups = self.num_attention_heads // self.num_key_value_heads
+        assert self.num_key_value_groups == 1, "Grouped linear not yet supported."
+
+        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.g_proj = nn.Linear(
+            args.hidden_size, args.num_attention_heads * self.head_dim, bias=False
+        )
+        self.g_norm = GroupRMSNorm(
+            args.num_attention_heads * self.head_dim,
+            eps=args.rms_norm_eps,
+            groups=args.group_norm_size,
+        )
+
+        if args.use_qk_norm:
+            self.key_layernorm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
+            self.query_layernorm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
+
+        self.rope = initialize_rope(
+            int(self.head_dim * args.partial_rotary_factor),
+            args.rope_theta,
+            traditional=args.rope_traditional,
+            scaling_config=args.rope_scaling,
+            max_position_embeddings=args.max_position_embeddings,
+        )
+        self._slope = self._get_slopes()
+
+    def _get_slopes(self) -> mx.array:
+        n = self.num_attention_heads
+
+        def power_of_2_slopes(n):
+            return [2 ** (-(2 ** -(math.log2(n) - 3)) * (i + 1)) for i in range(n)]
+
+        if math.log2(n).is_integer():
+            slopes = power_of_2_slopes(n)
+        else:
+            p = 2 ** math.floor(math.log2(n))
+            slopes = power_of_2_slopes(p) + power_of_2_slopes(2 * p)[::2][: n - p]
+
+        slopes = mx.array(slopes, dtype=mx.float32)
+        denom = max(1, self.num_hidden_layers - 1)
+        layer_pos = max(0, self.layer_idx - 1)
+        layer_factor = 1 - (layer_pos / denom) + 1e-5
+        return -slopes * layer_factor
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+        offset: int = 0,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        qkv = self.query_key_value(x)
+        qkv_mix = qkv.reshape(
+            B,
+            L,
+            (self.num_attention_heads + 2 * self.num_key_value_heads),
+            self.head_dim,
+        )
+        q, k, v = mx.split(
+            qkv_mix,
+            [
+                self.num_attention_heads,
+                self.num_attention_heads + self.num_key_value_heads,
+            ],
+            axis=2,
+        )
+
+        queries = q.transpose(0, 2, 1, 3)
+        keys = k.transpose(0, 2, 1, 3)
+        values = v.transpose(0, 2, 1, 3)
+
+        if self.use_qk_norm:
+            queries = self.query_layernorm(queries)
+            keys = self.key_layernorm(keys)
+
+        queries = self.rope(queries, offset=offset)
+        keys = self.rope(keys, offset=offset)
+
+        if cache is None:
+            cache = [None]
+        output, cache[0] = recurrent_gla(
+            q=queries,
+            k=keys,
+            v=values,
+            g=self._slope,
+            scale=self.scale,
+            h=cache[0],
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        output = self.g_norm(output) * mx.sigmoid(self.g_proj(x))
+        return self.dense(output)
+
+
+def group_expert_select(
+    gates: mx.array,
+    e_score_correction_bias: mx.array,
+    top_k: int,
+    n_group: int,
+    topk_group: int,
+    routed_scaling_factor: float,
+    norm_topk_prob: bool,
+    score_function: str,
+) -> Tuple[mx.array, mx.array]:
+    in_type = gates.dtype
+    if score_function == "sigmoid":
+        scores = mx.sigmoid(gates.astype(mx.float32))
+    else:
+        scores = mx.softmax(gates.astype(mx.float32), axis=-1)
+    orig_scores = scores
+    if e_score_correction_bias is not None:
+        scores = scores + e_score_correction_bias
+    if n_group > 1:
+        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.astype(in_type)
+
+
+class Gate(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.norm_topk_prob = args.norm_topk_prob
+
+        self.top_k = args.num_experts_per_tok
+        self.n_group = args.n_group
+        self.topk_group = args.topk_group
+        self.routed_scaling_factor = args.routed_scaling_factor
+        self.enable_routed_scaling = args.moe_router_enable_routed_scaling
+
+        self.gate_proj = nn.Linear(args.hidden_size, args.num_experts, bias=False)
+        self.expert_bias = (
+            mx.zeros((args.num_experts,))
+            if args.moe_router_enable_expert_bias
+            else None
+        )
+        self.score_function = args.score_function
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return group_expert_select(
+            self.gate_proj(x),
+            self.expert_bias,
+            self.top_k,
+            self.n_group,
+            self.topk_group,
+            self.routed_scaling_factor,
+            self.norm_topk_prob,
+            self.score_function,
+        )
+
+
+class SparseMoeBlock(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 = Gate(args)
+        shared_dim = (
+            args.moe_shared_expert_intermediate_size or args.moe_intermediate_size
+        )
+        self.shared_experts = (
+            MLP(
+                args=args,
+                intermediate_size=shared_dim * args.num_shared_experts,
+            )
+            if args.num_shared_experts > 0 and args.moe_router_enable_shared_expert
+            else None
+        )
+
+    def __call__(self, x: mx.array) -> mx.array:
+        topk_idx, topk_weight = self.gate(x)
+        out = self.switch_mlp(x, topk_idx)
+        out = (out * topk_weight[..., None]).sum(axis=-2)
+        if self.shared_experts is not None:
+            out = out + self.shared_experts(x)
+        return out
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.is_global = (
+            (layer_idx + 1) % args.layer_group_size == 0
+            or layer_idx
+            >= args.num_hidden_layers // args.layer_group_size * args.layer_group_size
+        )
+
+        if self.is_global:
+            self.attention = Attention(args)
+        else:
+            self.attention = LinearAttention(args, layer_idx=layer_idx)
+
+        self.mlp = (
+            SparseMoeBlock(args)
+            if (
+                args.num_experts is not None and layer_idx >= args.first_k_dense_replace
+            )
+            else 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,
+        offset: int = 0,
+    ) -> mx.array:
+        if self.is_global:
+            r = self.attention(self.input_layernorm(x), mask, cache)
+        else:
+            r = self.attention(self.input_layernorm(x), mask, cache, offset=offset)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        return h + r
+
+
+class LanguageModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.word_embeddings = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            DecoderLayer(args, layer_idx=i) for i in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.gla_idx = 0
+        self.attn_idx = args.layer_group_size - 1
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        h = self.word_embeddings(inputs)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        offset = 0
+        attn_mask = create_attention_mask(h, cache[self.attn_idx])
+        gla_mask = create_ssm_mask(h, cache[self.gla_idx])
+        if cache[self.attn_idx] is not None:
+            offset = cache[self.attn_idx].offset
+
+        for layer, c in zip(self.layers, cache):
+            mask = attn_mask if layer.is_global else gla_mask
+            h = layer(h, mask, c, offset=offset)
+
+        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 = LanguageModel(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        out = self.model(inputs, cache)
+        if self.args.tie_word_embeddings:
+            out = self.model.word_embeddings.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)
+
+        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}"
+            # Handle MoE layers
+            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.gate_proj"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate
+
+    @property
+    def cast_predicate(self):
+        def predicate(k):
+            return "expert_bias" not in k
+
+        return predicate
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def make_cache(self):
+        caches = []
+        for l in self.layers:
+            if l.is_global:
+                caches.append(KVCache())
+            else:
+                caches.append(ArraysCache(size=1))
+        return caches

mlx_lm/models/cache.py

@@ -1,5 +1,6 @@
 # Copyright © 2023-2024 Apple Inc.
 
+import copy
 from typing import Any, Dict, List, Optional
 
 import mlx.core as mx
@@ -73,10 +74,10 @@ def load_prompt_cache(file_name, return_metadata=False):
     arrays = tree_unflatten(list(arrays.items()))
     cache_metadata = tree_unflatten(list(cache_metadata.items()))
     info, metadata, classes = cache_metadata
-    cache = [globals()[c]() for c in classes]
-    for c, state, meta_state in zip(cache, arrays, info):
-        c.state = state
-        c.meta_state = meta_state
+    cache = [
+        globals()[c].from_state(state, meta_state)
+        for c, state, meta_state in zip(classes, arrays, info)
+    ]
     if return_metadata:
         return cache, metadata
     return cache
@@ -108,6 +109,10 @@ def trim_prompt_cache(cache: List[Any], num_tokens: int) -> List[Any]:
     return [c.trim(num_tokens) for c in cache][0]
 
 
+def cache_length(cache: List[Any]):
+    return max(len(c) for c in cache)
+
+
 def create_attention_mask(
     N: int, offset: int, return_array: bool, window_size: Optional[int]
 ):
@@ -141,6 +146,32 @@ class _BaseCache:
     def is_trimmable(self):
         return False
 
+    def __len__(self):
+        """The length of a cache is meant to represent the number of elements
+        that we need to process in the attention. For instance for KVCache it
+        is the size of the state, for RotatingKVCache it would be up to
+        max_size etc."""
+        return 0
+
+    def __bool__(self):
+        """When an object defines __len__ then python defines the bool operator
+        as len(obj) != 0. This, for instance, doesn't allow us to write
+
+            cache = cache or make_cache()
+
+        which is why we are overriding that behaviour with a constant bool
+        operator return True.
+        """
+        return True
+
+    @classmethod
+    def from_state(cls, state, meta_state):
+        # Create an instance of cls without calling __init__
+        obj = cls.__new__(cls)
+        obj.state = state
+        obj.meta_state = meta_state
+        return obj
+
 
 class ConcatenateKVCache(_BaseCache):
     """ConcatenateKVCache the simplest KV cache implementation.
@@ -188,11 +219,12 @@ class ConcatenateKVCache(_BaseCache):
 
 
 class QuantizedKVCache(_BaseCache):
+    step = 256
+
     def __init__(self, group_size: int = 64, bits: int = 8):
         self.keys = None
         self.values = None
         self.offset = 0
-        self.step = 256
         self.group_size = group_size
         self.bits = bits
 
@@ -254,11 +286,11 @@ class QuantizedKVCache(_BaseCache):
 
     @property
     def meta_state(self):
-        return tuple(map(str, (self.step, self.offset, self.group_size, self.bits)))
+        return tuple(map(str, (self.offset, self.group_size, self.bits)))
 
     @meta_state.setter
     def meta_state(self, v):
-        self.step, self.offset, self.group_size, self.bits = map(int, v)
+        self.offset, self.group_size, self.bits = map(int, v)
 
     def is_trimmable(self):
         return True
@@ -273,11 +305,12 @@ class QuantizedKVCache(_BaseCache):
 
 
 class KVCache(_BaseCache):
+    step = 256
+
     def __init__(self):
         self.keys = None
         self.values = None
         self.offset = 0
-        self.step = 256
 
     def update_and_fetch(self, keys, values):
         prev = self.offset
@@ -303,6 +336,9 @@ class KVCache(_BaseCache):
         self.values[..., prev : self.offset, :] = values
         return self.keys[..., : self.offset, :], self.values[..., : self.offset, :]
 
+    def __len__(self):
+        return self.offset
+
     @property
     def state(self):
         if self.offset == self.keys.shape[2]:
@@ -341,14 +377,14 @@ class KVCache(_BaseCache):
 
 
 class RotatingKVCache(_BaseCache):
+    step = 256
 
-    def __init__(self, max_size=None, keep=0, step=256):
+    def __init__(self, max_size, keep=0):
         self.keep = keep
         self.keys = None
         self.values = None
         self.offset = 0
         self.max_size = max_size
-        self.step = step
         self._idx = 0
 
     def _trim(self, trim_size, v, append=None):
@@ -388,10 +424,11 @@ class RotatingKVCache(_BaseCache):
             # preserve context
             self.keys = self._temporal_order(self.keys)
             self.values = self._temporal_order(self.values)
+            self._idx = self.keys.shape[2]
 
-            # The largest size is self.max_size + S to ensure
+            # The largest size is self.max_size + S - 1 to ensure
             # every token gets at least self.max_size context
-            trim_size = self._idx - self.max_size
+            trim_size = self._idx - self.max_size + 1
             self.keys = self._trim(trim_size, self.keys, keys)
             self.values = self._trim(trim_size, self.values, values)
         self.offset += keys.shape[2]
@@ -446,6 +483,9 @@ class RotatingKVCache(_BaseCache):
             return self._update_in_place(keys, values)
         return self._update_concat(keys, values)
 
+    def __len__(self):
+        return min(self.offset, self.max_size)
+
     @property
     def state(self):
         if self.offset < self.keys.shape[2]:
@@ -459,13 +499,11 @@ class RotatingKVCache(_BaseCache):
 
     @property
     def meta_state(self):
-        return tuple(
-            map(str, (self.keep, self.max_size, self.step, self.offset, self._idx))
-        )
+        return tuple(map(str, (self.keep, self.max_size, self.offset, self._idx)))
 
     @meta_state.setter
     def meta_state(self, v):
-        self.keep, self.max_size, self.step, self.offset, self._idx = map(
+        self.keep, self.max_size, self.offset, self._idx = map(
             int,
             v,
         )
@@ -487,7 +525,7 @@ class RotatingKVCache(_BaseCache):
     ):
         if N > 1:
             window_size = window_size or self.max_size
-            offset = min(self.max_size, self.offset)
+            offset = min(self.max_size - 1, self.offset)
             if offset + N > window_size or return_array:
                 return create_causal_mask(N, offset, window_size=window_size)
             else:
@@ -500,16 +538,19 @@ class RotatingKVCache(_BaseCache):
                 idx = self._idx
                 if idx >= self.max_size:
                     idx = 0
-                mask_size = min(self.max_size, self.offset)
+                if self.offset < self.max_size:
+                    mask_size = self.offset + 1
+                else:
+                    mask_size = self.max_size
                 mask = mx.arange(mask_size) >= (mask_size - window_size)
                 mask = mx.roll(mask, shift=idx + 1)
-                return mask[:, None]
+                return mask
 
 
 class ArraysCache(_BaseCache):
     def __init__(self, size, left_padding: Optional[List[int]] = None):
         self.cache = [None] * size
-        self.left_padding = left_padding
+        self.left_padding = mx.array(left_padding) if left_padding else None
 
     def __setitem__(self, idx, value):
         self.cache[idx] = value
@@ -552,7 +593,7 @@ class MambaCache(ArraysCache):
 
 
 class ChunkedKVCache(KVCache):
-    def __init__(self, chunk_size=None):
+    def __init__(self, chunk_size):
         super().__init__()
         self.chunk_size = chunk_size
         self.start_position = 0
@@ -603,7 +644,7 @@ class ChunkedKVCache(KVCache):
         self.chunk_size, self.start_position = map(int, v)
 
 
-class CacheList(KVCache):
+class CacheList(_BaseCache):
     def __init__(self, *caches):
         self.caches = caches
 
@@ -631,8 +672,33 @@ class CacheList(KVCache):
             c.state = v[start : start + l]
             start += l
 
+    def filter(self, batch_indices):
+        """
+        In-place filter to keep just the given indices in the cache.
+        """
+        for c in self.caches:
+            c.filter(batch_indices)
+
+    def extend(self, other):
+        """
+        In-place extend this cache with the other cache.
+        """
+        for c, o in zip(self.caches, other.caches):
+            c.extend(o)
+
+
+def dynamic_roll(x, shifts, axis):
+    n = x.shape[axis]
+    expand_shifts = (...,) + (None,) * (x.ndim - axis)
+    expand_indices = expand_shifts[:-1]
+    idx = (mx.arange(n)[expand_indices] - shifts[expand_shifts]) % n
+    rolled = mx.take_along_axis(x, idx, axis=axis)
+    return rolled
+
 
 class BatchKVCache(_BaseCache):
+    step = 256
+
     def __init__(self, left_padding: List[int]):
         """
         The BatchKV cache expects inputs to be left-padded.
@@ -657,7 +723,8 @@ class BatchKVCache(_BaseCache):
         self.left_padding = mx.array(left_padding)
         self.offset = mx.array([-l for l in left_padding])
         self._idx = 0
-        self.step = 256
+
+        self._right_padding = None
 
     def update_and_fetch(self, keys, values):
         prev = self._idx
@@ -684,6 +751,31 @@ class BatchKVCache(_BaseCache):
         self.values[..., prev : self._idx, :] = values
         return self.keys[..., : self._idx, :], self.values[..., : self._idx, :]
 
+    def __len__(self):
+        return self._idx
+
+    def prepare(self, *, left_padding=None, lengths=None, right_padding=None):
+        if left_padding is not None:
+            if self.keys is not None:
+                raise ValueError(
+                    "Left padding can only be added to an empty BatchKVCache"
+                )
+            left_padding = mx.array(left_padding)
+            self.left_padding += left_padding
+            self.offset -= left_padding
+
+        if right_padding is not None and max(right_padding) > 0:
+            self._right_padding = mx.array(right_padding)
+
+    def finalize(self):
+        if self._right_padding is not None:
+            padding = self._right_padding
+            self.keys = dynamic_roll(self.keys, padding[:, None], axis=2)
+            self.values = dynamic_roll(self.values, padding[:, None], axis=2)
+            self.offset -= padding
+            self.left_padding += padding
+            self._right_padding = None
+
     @property
     def state(self):
         k, v = self.keys, self.values
@@ -756,3 +848,351 @@ class BatchKVCache(_BaseCache):
             mx.concatenate, zip(*(pad(self), pad(other)))
         )
         self._idx = max_idx
+
+    def extract(self, idx):
+        cache = KVCache()
+        padding = self.left_padding[idx].item()
+        cache.keys = mx.contiguous(self.keys[idx : idx + 1, :, padding : self._idx])
+        cache.values = mx.contiguous(self.values[idx : idx + 1, :, padding : self._idx])
+        cache.offset = cache.keys.shape[2]
+        return cache
+
+    @classmethod
+    def merge(cls, caches):
+        lengths = [len(c) for c in caches]
+        max_length = max(lengths)
+        padding = [max_length - l for l in lengths]
+        B = len(caches)
+        H = max(c.keys.shape[1] for c in caches if c.keys is not None)
+        Dk = max(c.keys.shape[3] for c in caches if c.keys is not None)
+        Dv = max(c.values.shape[3] for c in caches if c.values is not None)
+        dt = next(iter(c.keys.dtype for c in caches if c.keys is not None))
+
+        keys = mx.zeros((B, H, max_length, Dk), dtype=dt)
+        values = mx.zeros((B, H, max_length, Dv), dtype=dt)
+        for i, (p, c) in enumerate(zip(padding, caches)):
+            keys[i : i + 1, :, p : p + c.offset] = c.keys[..., : c.offset, :]
+            values[i : i + 1, :, p : p + c.offset] = c.values[..., : c.offset, :]
+
+        cache = cls(padding)
+        cache.keys = keys
+        cache.values = values
+        cache.offset += keys.shape[2]
+        cache._idx = keys.shape[2]
+
+        return cache
+
+
+class BatchRotatingKVCache(_BaseCache):
+    step = 256
+
+    def __init__(self, max_size, left_padding: List[int]):
+        self.keys = None
+        self.values = None
+
+        self.left_padding = mx.array(left_padding)
+        self.offset = mx.array([-l for l in left_padding])
+
+        self.max_size = max_size
+        self._idx = 0
+        self._offset = 0
+        self.rotated = False
+
+        # Lengths for right_padded inputs to make sure that padding tokens do
+        # not evict valid tokens.
+        self._lengths = None
+
+    def _trim(self, trim_size, v, append=None):
+        if trim_size > 0:
+            v = v[..., trim_size:, :]
+        if append is not None:
+            return mx.concatenate([v, append], axis=2)
+        return v
+
+    def _temporal_order(self):
+        """
+        Rearrange the cache into temporal order.
+        """
+        if self.rotated:
+            self.keys = mx.roll(self.keys, -self._idx, axis=2)
+            self.values = mx.roll(self.values, -self._idx, axis=2)
+            self._idx = self.keys.shape[2]
+            self.rotated = False
+
+    def _update_concat(self, keys, values):
+        if self.keys is None:
+            self.keys = keys
+            self.values = values
+        else:
+            # Put the keys/values in temporal order to
+            # preserve context
+            self._temporal_order()
+
+            # Slice off the end if needed
+            if self.keys.shape[2] > self._idx:
+                self.keys = self.keys[..., : self._idx, :]
+                self.values = self.values[..., : self._idx, :]
+
+            # Roll right sequences that are padded to make sure that we don't
+            # trim valid cache entries
+            if self._lengths is not None:
+                roll = mx.maximum(0, self.offset - self._lengths)
+                self.keys = dynamic_roll(self.keys, roll[:, None], axis=2)
+                self.values = dynamic_roll(self.values, roll[:, None], axis=2)
+                self.left_padding += roll
+                self.offset -= roll
+
+            # The largest size is self.max_size + S - 1 to ensure
+            # every token gets at least self.max_size context
+            trim_size = self._idx - self.max_size + 1
+            if trim_size > 0:
+                self.left_padding -= trim_size
+            self.keys = self._trim(trim_size, self.keys, keys)
+            self.values = self._trim(trim_size, self.values, values)
+        self.offset += keys.shape[2]
+        self._offset += keys.shape[2]
+        self._idx = self.keys.shape[2]
+        return self.keys, self.values
+
+    def _update_in_place(self, keys, values):
+        if self._lengths is not None:
+            raise RuntimeError(
+                "finalize() should be called before deocoding with BatchRotatingKVCache"
+            )
+
+        # May not have hit the max size yet, so potentially
+        # keep growing the cache
+        B, n_kv_heads, S, k_head_dim = keys.shape
+        prev = self._offset
+        if self.keys is None or (
+            prev >= self.keys.shape[2] and self.keys.shape[2] < self.max_size
+        ):
+            v_head_dim = values.shape[3]
+            new_size = min(self.step, self.max_size - prev)
+            k_shape = (B, n_kv_heads, new_size, k_head_dim)
+            v_shape = (B, n_kv_heads, new_size, 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:
+                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._idx = prev
+
+        # Trim if needed
+        trim_size = self.keys.shape[2] - self.max_size
+        if trim_size > 0:
+            self.keys = self._trim(trim_size, self.keys)
+            self.values = self._trim(trim_size, self.values)
+            self._idx = self.max_size
+            self.left_padding -= trim_size
+
+        # Rotate
+        if self._idx == self.max_size:
+            self.rotated = True
+            self._idx = 0
+        if self.rotated:
+            self.left_padding -= S
+
+        # Assign
+        self.keys[..., self._idx : self._idx + S, :] = keys
+        self.values[..., self._idx : self._idx + S, :] = values
+        self._offset += S
+        self.offset += S
+        self._idx += S
+
+        # If the buffer is not full, slice off the end
+        if self._offset < self.max_size:
+            return (
+                self.keys[..., : self._offset, :],
+                self.values[..., : self._offset, :],
+            )
+        return self.keys, self.values
+
+    def update_and_fetch(self, keys, values):
+        if keys.shape[2] == 1:
+            return self._update_in_place(keys, values)
+        return self._update_concat(keys, values)
+
+    def __len__(self):
+        return min(self._offset, self.max_size)
+
+    def prepare(self, *, left_padding=None, lengths=None, right_padding=None):
+        if left_padding is not None:
+            if self.keys is not None:
+                raise ValueError(
+                    "Left padding can only be added to an empty BatchRotatingKVCache"
+                )
+            left_padding = mx.array(left_padding)
+            self.left_padding += left_padding
+            self.offset -= left_padding
+
+        if right_padding is not None and max(right_padding) > 0:
+            self._lengths = mx.array(lengths) + self.offset
+
+    def finalize(self):
+        if self._lengths is not None:
+            roll = mx.maximum(0, self.offset - self._lengths)
+            self.keys = dynamic_roll(self.keys, roll[:, None], axis=2)
+            self.values = dynamic_roll(self.values, roll[:, None], axis=2)
+            self.left_padding += roll
+            self.offset -= roll
+            self._lengths = None
+
+    @property
+    def state(self):
+        k, v = self.keys, self.values
+        if self._offset < k.shape[2]:
+            k, v = k[..., : self._offset, :], v[..., : self._offset, :]
+        return k, v, self.offset, self.left_padding
+
+    @state.setter
+    def state(self, v):
+        self.keys, self.values, self.offset, self.left_padding = v
+
+    @property
+    def meta_state(self):
+        return tuple(map(str, (self.max_size, self._offset, self._idx, self.rotated)))
+
+    @meta_state.setter
+    def meta_state(self, v):
+        self.max_size, self._offset, self._idx = map(
+            int,
+            v[:3],
+        )
+        self.rotated = bool(v[3])
+
+    def is_trimmable(self):
+        return self._offset < self.max_size
+
+    def trim(self, n):
+        n = min(self._offset, n)
+        self._offset -= n
+        self._idx -= n
+        self.offset -= n
+        return n
+
+    def to_quantized(self, group_size: int = 64, bits: int = 4) -> QuantizedKVCache:
+        raise NotImplementedError("BatchRotatingKVCache Quantization NYI")
+
+    def make_mask(
+        self, N: int, window_size: Optional[int] = None, return_array: bool = False
+    ):
+        left_padding = self.left_padding
+        window_size = window_size or self.max_size
+        offset = min(self.max_size - 1, self._offset)
+        rinds = mx.arange(offset + N)
+        linds = mx.arange(offset, offset + N) if offset else rinds
+        linds = linds[:, None]
+        rinds = rinds[None]
+        mask = linds >= rinds
+        mask &= linds < rinds + window_size
+        if (trim_size := self._idx - self.max_size + int(N > 1)) > 0:
+            left_padding = left_padding - trim_size
+
+        rotated = N == 1 and (self.rotated or self._idx >= self.max_size)
+        if rotated:
+            left_padding = left_padding - 1
+
+        mask = mask & (rinds >= mx.expand_dims(left_padding, (1, 2, 3)))
+
+        if rotated:
+            idx = self._idx
+            if idx >= self.max_size:
+                idx = 0
+            mask = mx.roll(mask, shift=idx + 1, axis=-1)
+
+        return mask
+
+    def filter(self, batch_indices):
+        """
+        In-place filter to keep just the given indices in the cache.
+        """
+        self.keys = self.keys[batch_indices]
+        self.values = self.values[batch_indices]
+        self.offset = self.offset[batch_indices]
+        self.left_padding = self.left_padding[batch_indices]
+
+    def extend(self, other):
+        """
+        In-place extend this cache with the other cache.
+        """
+        if (self.rotated != other.rotated) or self._idx != other._idx:
+            self._temporal_order()
+            other._temporal_order()
+
+        max_idx = max(self._idx, other._idx)
+        max_size = max(self.keys.shape[2], other.keys.shape[2])
+
+        def pad(c):
+            left = max_idx - c._idx
+            right = max_size - c.keys.shape[2] - left
+            k, v = c.keys, c.values
+            if right < 0:
+                k = k[..., :right, :]
+                v = v[..., :right, :]
+                right = 0
+            if left != 0 or right != 0:
+                pad = [(0, 0), (0, 0), (left, right), (0, 0)]
+                k = mx.pad(k, pad)
+                v = mx.pad(v, pad)
+            left_padding = c.left_padding + left
+            return k, v, c.offset, left_padding
+
+        self.keys, self.values, self.offset, self.left_padding = map(
+            mx.concatenate, zip(*(pad(self), pad(other)))
+        )
+        self._idx = max_idx
+        self._offset = max(self._offset, other._offset)
+
+    def extract(self, idx):
+        cache = RotatingKVCache(self.max_size)
+        padding = self.left_padding[idx].item()
+        offset = self.offset[idx].item()
+        cache.keys = self.keys[idx : idx + 1]
+        cache.values = self.values[idx : idx + 1]
+        cache._idx = self._idx
+        if self.rotated:
+            cache.keys = mx.roll(cache.keys, -self._idx, axis=2)
+            cache.values = mx.roll(cache.values, -self._idx, axis=2)
+            cache._idx = self.max_size
+        if padding > 0:
+            cache.keys = mx.contiguous(cache.keys[:, :, padding : cache._idx])
+            cache.values = mx.contiguous(cache.values[:, :, padding : cache._idx])
+        cache.offset = offset
+        cache._idx = cache.keys.shape[2]
+
+        return cache
+
+    @classmethod
+    def merge(cls, caches):
+        if not all(c.max_size == caches[0].max_size for c in caches):
+            raise ValueError(
+                "BatchRotatingKVCache can only merge caches with the same maximum size"
+            )
+
+        offsets = [c.offset for c in caches]
+        lengths = [len(c) for c in caches]
+        max_length = max(lengths)
+        padding = [max_length - l for l in lengths]
+        B = len(caches)
+        H = max(c.keys.shape[1] for c in caches if c.keys is not None)
+        Dk = max(c.keys.shape[3] for c in caches if c.keys is not None)
+        Dv = max(c.values.shape[3] for c in caches if c.values is not None)
+        dt = next(iter(c.keys.dtype for c in caches if c.keys is not None))
+
+        keys = mx.zeros((B, H, max_length, Dk), dtype=dt)
+        values = mx.zeros((B, H, max_length, Dv), dtype=dt)
+        for i, (p, c) in enumerate(zip(padding, caches)):
+            keys[i : i + 1, :, p : p + c.offset] = c._temporal_order(c.keys)
+            values[i : i + 1, :, p : p + c.offset] = c._temporal_order(c.values)
+
+        cache = cls(caches[0].max_size, padding)
+        cache.keys = keys
+        cache.values = values
+        cache.offset = mx.array(offsets)
+        cache._idx = keys.shape[2]
+        cache._offset = keys.shape[2]
+
+        return cache

mlx_lm/models/deepseek_v2.py

@@ -8,6 +8,7 @@ import mlx.core as mx
 import mlx.nn as nn
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .pipeline import PipelineMixin
 from .switch_layers import SwitchGLU
 
 
@@ -355,7 +356,7 @@ class DeepseekV2DecoderLayer(nn.Module):
         return out
 
 
-class DeepseekV2Model(nn.Module):
+class DeepseekV2Model(PipelineMixin, nn.Module):
     def __init__(self, config: ModelArgs):
         super().__init__()
         self.vocab_size = config.vocab_size
@@ -364,32 +365,8 @@ class DeepseekV2Model(nn.Module):
             DeepseekV2DecoderLayer(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)
 
-        self.pipeline_rank = 0
-        self.pipeline_size = 1
-
-    def pipeline(self, group):
-        # Split layers in reverse so rank=0 gets the last layers and
-        # rank=pipeline_size-1 gets the first
-        self.pipeline_rank = group.rank()
-        self.pipeline_size = group.size()
-        layers_per_rank = len(self.layers) // self.pipeline_size
-        extra = len(self.layers) - layers_per_rank * self.pipeline_size
-        if self.pipeline_rank < extra:
-            layers_per_rank += 1
-
-        self.start_idx = (self.pipeline_size - self.pipeline_rank - 1) * layers_per_rank
-        self.end_idx = self.start_idx + layers_per_rank
-        self.num_layers = layers_per_rank
-        self.layers = self.layers[: self.end_idx]
-        self.layers[: self.start_idx] = [None] * self.start_idx
-        self.num_layers = len(self.layers) - self.start_idx
-
     def __call__(
         self,
         x: mx.array,
@@ -401,19 +378,21 @@ class DeepseekV2Model(nn.Module):
         pipeline_size = self.pipeline_size
 
         if cache is None:
-            cache = [None] * self.num_layers
+            cache = [None] * len(self.pipeline_layers)
         mask = create_attention_mask(h, cache[0])
 
         # Receive from the previous process in the pipeline
         if pipeline_rank < pipeline_size - 1:
             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])
+        for l, c in zip(self.pipeline_layers, cache):
+            h = l(h, mask, cache=c)
 
         # Send to the next process in the pipeline
         if pipeline_rank != 0:
             h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
+            if cache[-1] is not None:
+                cache[-1].keys = mx.depends(cache[-1].keys, h)
 
         # Broadcast h while keeping it in the graph
         h = mx.distributed.all_gather(h)[: h.shape[0]]
@@ -452,4 +431,4 @@ class Model(nn.Module):
 
     @property
     def layers(self):
-        return self.model.layers[self.model.start_idx : self.model.end_idx]
+        return self.model.pipeline_layers

mlx_lm/models/deepseek_v3.py

@@ -9,6 +9,8 @@ import mlx.core as mx
 import mlx.nn as nn
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .pipeline import PipelineMixin
+from .rope_utils import initialize_rope
 from .switch_layers import SwitchGLU
 
 
@@ -45,85 +47,6 @@ class ModelArgs(BaseModelArgs):
     attention_bias: bool = False
 
 
-def yarn_find_correction_dim(
-    num_rotations, dim, base=10000, max_position_embeddings=2048
-):
-    return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (
-        2 * math.log(base)
-    )
-
-
-def yarn_find_correction_range(
-    low_rot, high_rot, dim, base=10000, max_position_embeddings=2048
-):
-    low = math.floor(
-        yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings)
-    )
-    high = math.ceil(
-        yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings)
-    )
-    return max(low, 0), min(high, dim - 1)
-
-
-def yarn_get_mscale(scale=1, mscale=1):
-    if scale <= 1:
-        return 1.0
-    return 0.1 * mscale * math.log(scale) + 1.0
-
-
-def yarn_linear_ramp_mask(min_val, max_val, dim):
-    if min_val == max_val:
-        max_val += 0.001  # Prevent singularity
-
-    linear_func = (mx.arange(dim, dtype=mx.float32) - min_val) / (max_val - min_val)
-    return mx.clip(linear_func, 0, 1)
-
-
-class DeepseekV3YarnRotaryEmbedding(nn.Module):
-    def __init__(
-        self,
-        dim,
-        max_position_embeddings=2048,
-        base=10000,
-        scaling_factor=1.0,
-        original_max_position_embeddings=4096,
-        beta_fast=32,
-        beta_slow=1,
-        mscale=1,
-        mscale_all_dim=0,
-    ):
-        super().__init__()
-        self.mscale = yarn_get_mscale(scaling_factor, mscale) / yarn_get_mscale(
-            scaling_factor, mscale_all_dim
-        )
-        freq_extra = 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,
-            dim,
-            base,
-            original_max_position_embeddings,
-        )
-        freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dim // 2)
-        self._freqs = (freq_inter * freq_extra) / (
-            freq_inter * freq_mask + freq_extra * (1 - freq_mask)
-        )
-
-    def __call__(self, x, offset=0):
-        if self.mscale != 1.0:
-            x = self.mscale * x
-        return mx.fast.rope(
-            x,
-            x.shape[-1],
-            traditional=True,
-            base=None,
-            scale=1.0,
-            offset=offset,
-            freqs=self._freqs,
-        )
-
-
 class DeepseekV3Attention(nn.Module):
     def __init__(self, config: ModelArgs):
         super().__init__()
@@ -175,35 +98,19 @@ class DeepseekV3Attention(nn.Module):
 
         if self.config.rope_scaling is not None:
             mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0)
-            scaling_factor = self.config.rope_scaling["factor"]
             if mscale_all_dim:
-                mscale = yarn_get_mscale(scaling_factor, mscale_all_dim)
-                self.scale = self.scale * mscale * mscale
+                scaling_factor = self.config.rope_scaling["factor"]
+                if scaling_factor > 1:
+                    s = 0.1 * mscale_all_dim * math.log(scaling_factor) + 1.0
+                    self.scale = self.scale * s * s
 
-            rope_kwargs = {
-                key: self.config.rope_scaling[key]
-                for key in [
-                    "original_max_position_embeddings",
-                    "beta_fast",
-                    "beta_slow",
-                    "mscale",
-                    "mscale_all_dim",
-                ]
-                if key in self.config.rope_scaling
-            }
-            self.rope = DeepseekV3YarnRotaryEmbedding(
-                dim=self.qk_rope_head_dim,
-                max_position_embeddings=self.max_position_embeddings,
-                scaling_factor=scaling_factor,
-                base=self.rope_theta,
-                **rope_kwargs,
-            )
-        else:
-            self.rope = nn.RoPE(
-                dims=self.qk_rope_head_dim,
-                base=self.rope_theta,
-                traditional=True,
-            )
+        self.rope = initialize_rope(
+            dims=self.qk_rope_head_dim,
+            base=self.rope_theta,
+            traditional=True,
+            max_position_embeddings=self.max_position_embeddings,
+            scaling_config=self.config.rope_scaling,
+        )
 
     def __call__(
         self,
@@ -389,7 +296,7 @@ class DeepseekV3DecoderLayer(nn.Module):
         return h + r
 
 
-class DeepseekV3Model(nn.Module):
+class DeepseekV3Model(PipelineMixin, nn.Module):
     def __init__(self, config: ModelArgs):
         super().__init__()
         self.vocab_size = config.vocab_size
@@ -398,28 +305,7 @@ class DeepseekV3Model(nn.Module):
             DeepseekV3DecoderLayer(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)
-        self.pipeline_rank = 0
-        self.pipeline_size = 1
-
-    def pipeline(self, group):
-        # Split layers in reverse so rank=0 gets the last layers and
-        # rank=pipeline_size-1 gets the first
-        self.pipeline_rank = group.rank()
-        self.pipeline_size = group.size()
-        layers_per_rank = len(self.layers) // self.pipeline_size
-        extra = len(self.layers) - layers_per_rank * self.pipeline_size
-        if self.pipeline_rank < extra:
-            layers_per_rank += 1
-        self.start_idx = (self.pipeline_size - self.pipeline_rank - 1) * layers_per_rank
-        self.end_idx = self.start_idx + layers_per_rank
-        self.layers = self.layers[: self.end_idx]
-        self.layers[: self.start_idx] = [None] * self.start_idx
-        self.num_layers = len(self.layers) - self.start_idx
 
     def __call__(
         self,
@@ -432,20 +318,21 @@ class DeepseekV3Model(nn.Module):
         pipeline_size = self.pipeline_size
 
         if cache is None:
-            cache = [None] * self.num_layers
+            cache = [None] * len(self.pipeline_layers)
         mask = create_attention_mask(h, cache[0])
 
         # Receive from the previous process in the pipeline
-
         if pipeline_rank < pipeline_size - 1:
             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])
+        for l, c in zip(self.pipeline_layers, cache):
+            h = l(h, mask, cache=c)
 
         # Send to the next process in the pipeline
         if pipeline_rank != 0:
             h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
+            if cache[-1] is not None:
+                cache[-1].keys = mx.depends(cache[-1].keys, h)
 
         # Broadcast h while keeping it in the graph
         h = mx.distributed.all_gather(h)[: h.shape[0]]
@@ -485,7 +372,22 @@ class Model(nn.Module):
             )
             return weight[:m, :n].astype(dtype)
 
-        # Dequantize
+        # Remap for int4
+        new_weights = {}
+        for k, v in weights.items():
+            if k.endswith("weight_shape"):
+                base = k.replace("weight_shape", "")
+                new_weights[base + "weight"] = weights[base + "weight_packed"].view(
+                    mx.uint32
+                )
+                s = weights[base + "weight_scale"]
+                new_weights[base + "scales"] = s
+                new_weights[base + "biases"] = -8 * s
+            elif not (k.endswith("weight_scale") or k.endswith("weight_packed")):
+                new_weights[k] = v
+        weights = new_weights
+
+        # Dequantize fp8
         new_weights = {}
         for k, v in weights.items():
             if "weight_scale_inv" in k:
@@ -519,7 +421,7 @@ class Model(nn.Module):
 
     @property
     def layers(self):
-        return self.model.layers[self.model.start_idx : self.model.end_idx]
+        return self.model.pipeline_layers
 
     @property
     def cast_predicate(self):

mlx_lm/models/deepseek_v32.py

@@ -0,0 +1,515 @@
+# Copyright © 2025 Apple Inc.
+
+import math
+from dataclasses import dataclass
+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 .cache import CacheList, KVCache
+from .rope_utils import initialize_rope
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str = "deepseek_v32"
+    vocab_size: int = 102400
+    hidden_size: int = 4096
+    index_head_dim: int = 128
+    index_n_heads: int = 64
+    index_topk: int = 2048
+    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: int = 1
+    topk_group: int = 1
+    num_experts_per_tok: int = 1
+    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
+
+
+class Indexer(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.dim = args.hidden_size
+        self.n_heads = args.index_n_heads
+        self.head_dim = args.index_head_dim
+        self.rope_head_dim = args.qk_rope_head_dim
+        self.index_topk = args.index_topk
+        self.q_lora_rank = args.q_lora_rank
+        self.wq_b = nn.Linear(
+            self.q_lora_rank, self.n_heads * self.head_dim, bias=False
+        )
+        self.wk = nn.Linear(self.dim, self.head_dim, bias=False)
+        self.k_norm = nn.LayerNorm(self.head_dim)
+        self.weights_proj = nn.Linear(self.dim, self.n_heads, bias=False)
+        self.softmax_scale = self.head_dim**-0.5
+        self.rope = initialize_rope(
+            dims=args.qk_rope_head_dim,
+            base=args.rope_theta,
+            traditional=False,
+            max_position_embeddings=args.max_position_embeddings,
+            scaling_config=args.rope_scaling,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        qr: mx.array,
+        mask: Optional[mx.array],
+        cache: Optional[Any] = None,
+    ):
+        # Computes top_k indices for attention
+        b, s, _ = x.shape
+        q = self.wq_b(qr)
+        q = q.reshape(b, s, self.n_heads, self.head_dim).swapaxes(1, 2)
+        q_pe, q_nope = mx.split(q, [self.rope_head_dim], axis=-1)
+
+        offset = cache.offset if cache is not None else 0
+
+        q_pe = self.rope(q_pe, offset=offset)
+        q = mx.concatenate([q_pe, q_nope], axis=-1)
+
+        k = self.wk(x)
+        k = self.k_norm(k)
+        k = mx.reshape(k, (b, 1, s, self.head_dim))
+        k_pe, k_nope = mx.split(k, [self.rope_head_dim], axis=-1)
+        k_pe = self.rope(k_pe, offset=offset)
+        k = mx.concatenate([k_pe, k_nope], axis=-1)
+        if cache is not None:
+            k, _ = cache.update_and_fetch(k, mx.zeros([b, 1, s, 0]))
+        if k.shape[2] <= self.index_topk:
+            return None
+        scores = q @ k.swapaxes(-1, -2)
+        scores = mx.maximum(scores, 0)
+        weights = self.weights_proj(x) * (self.n_heads**-0.5 * self.softmax_scale)
+        weights = weights.swapaxes(-1, -2)[..., None]
+        scores = scores * weights
+        scores = scores.sum(axis=1)
+        if mask is not None:
+            scores = mx.where(mask, scores, -float("inf"))
+        return mx.argpartition(scores, kth=-self.index_topk, axis=-1)[
+            ..., -self.index_topk :
+        ]
+
+
+class DeepseekV32Attention(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rope_theta = config.rope_theta
+        self.q_lora_rank = config.q_lora_rank
+        self.qk_rope_head_dim = config.qk_rope_head_dim
+        self.kv_lora_rank = config.kv_lora_rank
+        self.v_head_dim = config.v_head_dim
+        self.qk_nope_head_dim = config.qk_nope_head_dim
+        self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim
+
+        self.scale = self.q_head_dim**-0.5
+
+        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, eps=1e-6)
+        self.q_b_proj = nn.Linear(
+            self.q_lora_rank, self.num_heads * self.q_head_dim, bias=False
+        )
+
+        self.kv_a_proj_with_mqa = nn.Linear(
+            self.hidden_size,
+            self.kv_lora_rank + self.qk_rope_head_dim,
+            bias=config.attention_bias,
+        )
+        self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank, eps=1e-6)
+        self.kv_b_proj = nn.Linear(
+            self.kv_lora_rank,
+            self.num_heads
+            * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
+            bias=False,
+        )
+
+        self.o_proj = nn.Linear(
+            self.num_heads * self.v_head_dim,
+            self.hidden_size,
+            bias=config.attention_bias,
+        )
+
+        if self.config.rope_scaling is not None:
+            mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0)
+            if mscale_all_dim:
+                scaling_factor = self.config.rope_scaling["factor"]
+                if scaling_factor > 1:
+                    s = 0.1 * mscale_all_dim * math.log(scaling_factor) + 1.0
+                    self.scale = self.scale * s * s
+
+        self.indexer = Indexer(config)
+        self.rope = initialize_rope(
+            dims=self.qk_rope_head_dim,
+            base=self.rope_theta,
+            traditional=True,
+            max_position_embeddings=self.max_position_embeddings,
+            scaling_config=self.config.rope_scaling,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        qr = self.q_a_layernorm(self.q_a_proj(x))
+        q = self.q_b_proj(qr)
+
+        q = q.reshape(B, L, self.num_heads, self.q_head_dim).transpose(0, 2, 1, 3)
+        q_nope, q_pe = mx.split(q, [self.qk_nope_head_dim], axis=-1)
+        compressed_kv = self.kv_a_proj_with_mqa(x)
+        compressed_kv, k_pe = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
+        k_pe = k_pe.reshape(B, L, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3)
+        kv = self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
+        kv = kv.reshape(B, L, self.num_heads, -1).transpose(0, 2, 1, 3)
+
+        k_nope, values = mx.split(kv, [self.qk_nope_head_dim], axis=-1)
+
+        if cache is not None:
+            q_pe = self.rope(q_pe, cache[0].offset)
+            k_pe = self.rope(k_pe, cache[0].offset)
+            k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
+            keys, values = cache[0].update_and_fetch(
+                mx.concatenate([k_nope, k_pe], axis=-1), values
+            )
+        else:
+            cache = [None] * 2
+            q_pe = self.rope(q_pe)
+            k_pe = self.rope(k_pe)
+            k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
+            keys = mx.concatenate([k_nope, k_pe], axis=-1)
+
+        queries = mx.concatenate([q_nope, q_pe], axis=-1)
+        topk_indices = self.indexer(x, qr, mask, cache=cache[1])
+        if topk_indices is not None:
+            k_seq = keys.shape[2]
+            sparse_mask = mx.zeros((B, L, k_seq), dtype=mx.bool_)
+            sparse_mask = mx.put_along_axis(
+                sparse_mask, topk_indices, mx.array(True), axis=-1
+            )
+            sparse_mask = sparse_mask[:, None, :, :]
+            if mask is not None:
+                sparse_mask = sparse_mask & mask
+            mask = sparse_mask
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache[0], scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class DeepseekV32MLP(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:
+        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 DeepseekV32MoE(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 = DeepseekV32MLP(
+                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 DeepseekV32DecoderLayer(nn.Module):
+    def __init__(self, config: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.self_attn = DeepseekV32Attention(config)
+        self.mlp = (
+            DeepseekV32MoE(config)
+            if (
+                config.n_routed_experts is not None
+                and layer_idx >= config.first_k_dense_replace
+                and layer_idx % config.moe_layer_freq == 0
+            )
+            else DeepseekV32MLP(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 DeepseekV32Model(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 = [
+            DeepseekV32DecoderLayer(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)
+        self.pipeline_rank = 0
+        self.pipeline_size = 1
+
+    def pipeline(self, group):
+        # Split layers in reverse so rank=0 gets the last layers and
+        # rank=pipeline_size-1 gets the first
+        self.pipeline_rank = group.rank()
+        self.pipeline_size = group.size()
+        layers_per_rank = len(self.layers) // self.pipeline_size
+        extra = len(self.layers) - layers_per_rank * self.pipeline_size
+        if self.pipeline_rank < extra:
+            layers_per_rank += 1
+        self.start_idx = (self.pipeline_size - self.pipeline_rank - 1) * layers_per_rank
+        self.end_idx = self.start_idx + layers_per_rank
+        self.layers = self.layers[: self.end_idx]
+        self.layers[: self.start_idx] = [None] * self.start_idx
+        self.num_layers = len(self.layers) - self.start_idx
+
+    def __call__(
+        self,
+        x: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        h = self.embed_tokens(x)
+
+        pipeline_rank = self.pipeline_rank
+        pipeline_size = self.pipeline_size
+
+        if cache is None:
+            cache = [None] * self.num_layers
+        mask = create_attention_mask(
+            h, cache[0][0] if cache[0] else None, return_array=True
+        )
+
+        # Receive from the previous process in the pipeline
+
+        if pipeline_rank < pipeline_size - 1:
+            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)
+            if cache[-1] is not None:
+                cache[-1].keys = mx.depends(cache[-1].keys, h)
+
+        # Broadcast h while keeping it in the graph
+        h = mx.distributed.all_gather(h)[: h.shape[0]]
+
+        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 = DeepseekV32Model(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ):
+        out = self.model(inputs, cache)
+        return self.lm_head(out)
+
+    def sanitize(self, weights):
+        def dequant(weight, scale_inv):
+            dtype = weight.dtype
+            bs = 128  # block size
+            m, n = weight.shape
+            pad_bottom = (-m) % bs
+            pad_side = (-n) % bs
+            weight = mx.pad(weight, ((0, pad_bottom), (0, pad_side)))
+            weight = weight.reshape(
+                ((m + pad_bottom) // bs, bs, (n + pad_side) // bs, bs)
+            )
+            weight = (weight * scale_inv[:, None, :, None]).reshape(
+                m + pad_bottom, n + pad_side
+            )
+            return weight[:m, :n].astype(dtype)
+
+        # Dequantize
+        new_weights = {}
+        for k, v in weights.items():
+            if "weight_scale_inv" in k:
+                scale_inv = v
+                wk = k.replace("_scale_inv", "")
+                weight = weights[wk]
+                weight = dequant(weight, scale_inv)
+                new_weights[wk] = weight
+            elif k not in new_weights:
+                new_weights[k] = v
+        weights = new_weights
+
+        # 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 and any unused precomputed rotary freqs
+        return {
+            k: v
+            for k, v in weights.items()
+            if not k.startswith("model.layers.61") and "rotary_emb.inv_freq" not in k
+        }
+
+    @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
+
+    def make_cache(self):
+        return [CacheList(KVCache(), KVCache()) for _ in self.layers]

mlx_lm/models/falcon_h1.py

@@ -0,0 +1,479 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass, field
+from typing import List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import (
+    BaseModelArgs,
+    create_attention_mask,
+    create_ssm_mask,
+    scaled_dot_product_attention,
+)
+from .cache import CacheList, KVCache, MambaCache
+from .rope_utils import initialize_rope
+from .ssm import ssm_update
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    attention_bias: bool = False
+    attention_in_multiplier: float = 1.0
+    attention_out_multiplier: float = 0.9375
+    embedding_multiplier: float = 5.656854249492381
+    head_dim: int = 64
+    hidden_size: int = 1024
+    initializer_range: float = 0.02
+    intermediate_size: int = 2048
+    key_multiplier: float = 0.390625
+    lm_head_multiplier: float = 0.0390625
+    mamba_chunk_size: int = 128
+    mamba_conv_bias: bool = True
+    mamba_d_conv: int = 4
+    mamba_d_head: int = 64
+    mamba_d_ssm: int = 1536
+    mamba_d_state: int = 128
+    mamba_expand: int = 2
+    mamba_n_groups: int = 1
+    mamba_n_heads: int = 24
+    mamba_norm_before_gate: bool = False
+    mamba_proj_bias: bool = False
+    mamba_rms_norm: bool = False
+    mamba_use_mlp: bool = True
+    max_position_embeddings: int = 131072
+    mlp_bias: bool = False
+    mlp_expansion_factor: int = 8
+    mlp_multipliers: List[float] = field(
+        default_factory=lambda: [0.8838834764831844, 0.5859375]
+    )
+    model_type: str = "falcon_h1"
+    num_attention_heads: int = 8
+    num_hidden_layers: int = 36
+    num_key_value_heads: int = 2
+    projectors_bias: bool = False
+    rms_norm_eps: float = 1e-05
+    rope_traditional: bool = False
+    rope_scaling: Optional[float] = None
+    rope_theta: float = 100000000000.0
+    ssm_in_multiplier: float = 1.25
+    ssm_multipliers: List[float] = field(
+        default_factory=lambda: [
+            0.3535533905932738,
+            0.25,
+            0.3535533905932738,
+            0.5,
+            0.3535533905932738,
+        ]
+    )
+    ssm_out_multiplier: float = 0.23570226039551587
+    vocab_size: int = 32784
+
+
+class FalconH1RMSNormGated(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6, n_groups=1, norm_before_gate=True):
+        super().__init__()
+        self.weight = mx.ones((hidden_size,))
+        self.variance_epsilon = eps
+        self.n_groups = n_groups
+        self.norm_before_gate = norm_before_gate
+
+    def __call__(self, hidden_states, gate=None):
+        if not self.norm_before_gate and gate is not None:
+            hidden_states = hidden_states * nn.silu(gate)
+
+        hidden_states = mx.fast.rms_norm(
+            hidden_states, self.weight, self.variance_epsilon
+        )
+
+        if self.norm_before_gate and gate is not None:
+            hidden_states = hidden_states * nn.silu(gate)
+        return hidden_states
+
+
+def compute_mup_vector(args):
+    intermediate_size = args.mamba_d_ssm
+    groups_time_state_size = args.mamba_n_groups * args.mamba_d_state
+    num_heads = args.mamba_n_heads
+    sizes = [
+        intermediate_size,
+        intermediate_size,
+        groups_time_state_size,
+        groups_time_state_size,
+        num_heads,
+    ]
+    return mx.concatenate(
+        [
+            mx.broadcast_to(mx.array(m), (s,))
+            for s, m in zip(sizes, args.ssm_multipliers)
+        ]
+    )
+
+
+class FalconH1Attention(nn.Module):
+
+    def __init__(self, args):
+        super().__init__()
+
+        self.hidden_size = args.hidden_size
+        self.num_heads = args.num_attention_heads
+        self.num_kv_heads = args.num_key_value_heads
+        self.head_dim = args.head_dim
+        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_kv_heads * self.head_dim,
+            bias=args.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            self.hidden_size,
+            self.num_kv_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
+        )
+
+        self.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            args.rope_traditional,
+            args.rope_scaling,
+            args.max_position_embeddings,
+        )
+
+    def __call__(self, x, mask=None, cache=None):
+        B, L, _ = x.shape
+
+        queries = self.q_proj(x)
+        keys = self.k_proj(x)
+        values = self.v_proj(x)
+
+        queries = queries.reshape(B, L, self.num_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, L, self.num_kv_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.num_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, mask=mask, scale=self.scale, cache=cache
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+
+        return self.o_proj(output)
+
+
+class FalconH1Mixer(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.num_heads = args.mamba_n_heads
+        self.hidden_size = args.hidden_size
+        self.ssm_state_size = args.mamba_d_state
+        self.conv_kernel_size = args.mamba_d_conv
+        self.intermediate_size = args.mamba_d_ssm
+        self.use_conv_bias = args.mamba_conv_bias
+
+        self.layer_norm_epsilon = args.rms_norm_eps
+        self.groups_time_state_size = args.mamba_n_groups * self.ssm_state_size
+
+        self.n_groups = args.mamba_n_groups
+        self.head_dim = args.mamba_d_head
+        self.chunk_size = args.mamba_chunk_size
+
+        self.time_step_limit = (0.0, float("inf"))
+        self.time_step_min = 0.001
+        self.time_step_max = 0.1
+
+        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,
+            bias=self.use_conv_bias,
+            kernel_size=self.conv_kernel_size,
+            groups=self.conv_dim,
+        )
+
+        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)
+
+        A = mx.arange(1, self.num_heads + 1)
+        self.A_log = mx.log(A)
+
+        self.mamba_rms_norm = args.mamba_rms_norm
+        if self.mamba_rms_norm:
+            self.norm = FalconH1RMSNormGated(
+                self.intermediate_size,
+                eps=self.layer_norm_epsilon,
+                n_groups=self.n_groups,
+                norm_before_gate=args.mamba_norm_before_gate,
+            )
+
+        self.D = mx.ones(self.num_heads)
+
+        self.out_proj = nn.Linear(
+            self.intermediate_size, self.hidden_size, bias=args.projectors_bias
+        )
+
+    def _apply_conv(
+        self, conv_input: mx.array, cache: Optional[MambaCache] = None
+    ) -> mx.array:
+        if cache is None or 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)
+
+        if cache is not None:
+            cache[0] = padded_input[:, -(self.conv_kernel_size - 1) :]
+
+        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,
+        state: Optional[mx.array] = None,
+        mask: Optional[mx.array] = None,
+    ) -> mx.array:
+        batch_size, seq_len, _ = hidden_states.shape
+
+        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)
+        C = C.reshape(batch_size, seq_len, self.n_groups, self.ssm_state_size)
+
+        y, state = ssm_update(
+            hidden_states,
+            self.A_log,
+            B,
+            C,
+            self.D,
+            dt,
+            self.dt_bias,
+            state,
+            self.time_step_limit,
+            mask,
+        )
+
+        return y.reshape(batch_size, seq_len, self.intermediate_size), state
+
+    def __call__(self, input_states, cache=None, mask: Optional[mx.array] = None):
+        projected_states = self.in_proj(input_states)
+
+        gate, conv_input, dt = mx.split(
+            projected_states,
+            [self.intermediate_size, self.intermediate_size + self.conv_dim],
+            axis=-1,
+        )
+
+        if mask is not None:
+            conv_input = mx.where(mask[..., None], conv_input, 0)
+        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,
+        )
+        state = cache[1] if cache else None
+        y, state = self._ssm(hidden_states_ssm, B, C, dt, state, mask)
+        if cache:
+            cache[1] = state
+
+        if self.mamba_rms_norm:
+            y = self.norm(y, gate)
+        else:
+            y = y * nn.silu(gate)
+
+        return self.out_proj(y)
+
+
+class FalconH1MLP(nn.Module):
+
+    def __init__(self, args):
+        super().__init__()
+
+        hidden_size = args.hidden_size
+        intermediate_size = args.intermediate_size
+
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=args.mlp_bias)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=args.mlp_bias)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=args.mlp_bias)
+
+    def __call__(self, x):
+        y = self.up_proj(x) * nn.silu(self.gate_proj(x))
+        y = self.down_proj(y)
+        return y
+
+
+class FalconH1DecoderLayer(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.feed_forward = FalconH1MLP(args)
+
+        head_dim = args.head_dim
+        self.channels_attn = (
+            args.num_attention_heads * head_dim
+            + 2 * args.num_key_value_heads * head_dim
+        )
+
+        self.mamba = FalconH1Mixer(args=args)
+
+        self.self_attn = FalconH1Attention(args)
+
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.pre_ff_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        h: mx.array,
+        cache,
+        attn_mask: Optional[mx.array],
+        mamba_mask: Optional[mx.array],
+    ) -> mx.array:
+
+        residual = h
+        h = self.input_layernorm(h)
+
+        mamba_h = self.mamba(input_states=h, cache=cache[0], mask=mamba_mask)
+
+        attn_h = self.self_attn(
+            h,
+            mask=attn_mask,
+            cache=cache[1],
+        )
+
+        h = residual + mamba_h + attn_h
+
+        residual = h
+        h = self.pre_ff_layernorm(h)
+        h = self.feed_forward(h)
+        return residual + h
+
+
+class FalconH1Model(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.hidden_size = args.hidden_size
+
+        self.embed_tokens = nn.Embedding(self.vocab_size, self.hidden_size)
+
+        self._mup_vector = compute_mup_vector(args)
+        self.layers = [
+            FalconH1DecoderLayer(args) for _ in range(args.num_hidden_layers)
+        ]
+        self.final_layernorm = nn.RMSNorm(self.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(self, inputs, cache=None):
+
+        h = self.embed_tokens(inputs)
+
+        h = h
+
+        if cache is None:
+            cache = [(None, None) * len(self.layers)]
+
+        mamba_mask = create_ssm_mask(h, cache[0][0])
+        attn_mask = create_attention_mask(h, cache[0][1])
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(
+                h,
+                cache=c,
+                attn_mask=attn_mask,
+                mamba_mask=mamba_mask,
+            )
+
+        return self.final_layernorm(h)
+
+
+class Model(nn.Module):
+
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = FalconH1Model(args=args)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(self, inputs, cache=None):
+        hidden_states = self.model(inputs, cache=cache)
+        return self.lm_head(hidden_states)
+
+    def sanitize(self, weights):
+        # Check if needs sanitization
+        c1d = weights["model.layers.0.mamba.conv1d.weight"]
+        if c1d.shape[-1] <= c1d.shape[1]:
+            return weights
+
+        sanitized_weights = {}
+        args = self.args
+
+        for name, param in weights.items():
+            # Fold-in multipliers
+            if name.endswith("embed_tokens.weight"):
+                param *= args.embedding_multiplier
+            elif name.endswith("lm_head.weight"):
+                param *= args.lm_head_multiplier
+            elif name.endswith("q_proj.weight") or name.endswith("k_proj.weight"):
+                param *= args.attention_in_multiplier
+            elif name.endswith("key_proj.weight"):
+                param *= args.attention_in_multiplier * args.key_multiplier
+            elif name.endswith("o_proj.weight"):
+                param *= args.attention_out_multiplier
+            elif name.endswith("out_proj.weight"):
+                param *= args.ssm_out_multiplier
+            elif name.endswith("gate_proj.weight"):
+                param *= args.mlp_multipliers[0]
+            elif name.endswith("down_proj.weight"):
+                param *= args.mlp_multipliers[1]
+            elif name.endswith("in_proj.weight"):
+                param *= (
+                    args.ssm_in_multiplier
+                    * self.model._mup_vector.astype(param.dtype)[:, None]
+                )
+            elif "conv1d.weight" in name:
+                param = param.transpose(0, 2, 1)
+            sanitized_weights[name] = param
+        return sanitized_weights
+
+    def make_cache(self):
+        return [
+            CacheList(MambaCache(), KVCache())
+            for _ in range(self.args.num_hidden_layers)
+        ]
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/gated_delta.py

@@ -7,15 +7,29 @@ import mlx.nn as nn
 
 @partial(mx.compile, shapeless=True)
 def compute_g(A_log, a, dt_bias):
-    return mx.exp(
-        -mx.exp(A_log.astype(mx.float32)) * nn.softplus(a + dt_bias).astype(A_log.dtype)
+    return mx.exp(-mx.exp(A_log.astype(mx.float32)) * nn.softplus(a + dt_bias)).astype(
+        A_log.dtype
     )
 
 
-def _make_gated_delta_kernel():
+def _make_gated_delta_kernel(has_mask=False, vectorized=False):
     if not mx.metal.is_available():
         return None
-    source = """
+    mask_source = "mask[b_idx * T + t]" if has_mask else "true"
+
+    # Configure g indexing based on whether gating is vectorized
+    if vectorized:
+        g_comment = "// g: [B, T, Hv, Dk]"
+        g_setup = "auto g_ = g + (b_idx * T * Hv + hv_idx) * Dk;"
+        g_access = "g_[s_idx]"
+        g_advance = "g_ += Hv * Dk;"
+    else:
+        g_comment = "// g: [B, T, Hv]"
+        g_setup = "auto g_ = g + b_idx * T * Hv;"
+        g_access = "g_[hv_idx]"
+        g_advance = "g_ += Hv;"
+
+    source = f"""
         auto n = thread_position_in_grid.z;
         auto b_idx = n / Hv;
         auto hv_idx = n % Hv;
@@ -38,60 +52,78 @@ def _make_gated_delta_kernel():
         auto o_state = state_out + (n * Dv + dv_idx) * Dk;
 
         float state[n_per_t];
-        for (int i = 0; i < n_per_t; ++i) {
-           auto s_idx = n_per_t * dk_idx + i;
-           state[i] = static_cast<float>(i_state[s_idx]);
-        }
+        for (int i = 0; i < n_per_t; ++i) {{
+          auto s_idx = n_per_t * dk_idx + i;
+          state[i] = static_cast<float>(i_state[s_idx]);
+        }}
 
-        // beta, g: [B, T, Hv]
-        auto g_ = g + b_idx * T * Hv;
+        {g_comment}
+        {g_setup}
         auto beta_ = beta + b_idx * T * Hv;
 
-        for (int t = 0; t < T; ++t) {
-          float kv_mem = 0.0f;
-          for (int i = 0; i < n_per_t; ++i) {
-            auto s_idx = n_per_t * dk_idx + i;
-            state[i] = state[i] * g_[hv_idx];
-            kv_mem += state[i] * k_[s_idx];
-          }
-          kv_mem = simd_sum(kv_mem);
+        for (int t = 0; t < T; ++t) {{
+          if ({mask_source}) {{
+            float kv_mem = 0.0f;
+            for (int i = 0; i < n_per_t; ++i) {{
+              auto s_idx = n_per_t * dk_idx + i;
+              state[i] = state[i] * {g_access};
+              kv_mem += state[i] * k_[s_idx];
+            }}
+            kv_mem = simd_sum(kv_mem);
 
-          auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx];
+            auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx];
 
-          float out = 0.0f;
-          for (int i = 0; i < n_per_t; ++i) {
-            auto s_idx = n_per_t * dk_idx + i;
-            state[i] = state[i] + k_[s_idx] * delta;
-            out += state[i] * q_[s_idx];
-          }
-          out = simd_sum(out);
-          if (thread_index_in_simdgroup == 0) {
-            y[dv_idx] = static_cast<InT>(out);
-          }
+            float out = 0.0f;
+            for (int i = 0; i < n_per_t; ++i) {{
+              auto s_idx = n_per_t * dk_idx + i;
+              state[i] = state[i] + k_[s_idx] * delta;
+              out += state[i] * q_[s_idx];
+            }}
+            out = simd_sum(out);
+            if (thread_index_in_simdgroup == 0) {{
+              y[dv_idx] = static_cast<InT>(out);
+            }}
+          }}
           // Increment data pointers to next time step
           q_ += Hk * Dk;
           k_ += Hk * Dk;
           v_ += Hv * Dv;
           y += Hv * Dv;
-          g_ += Hv;
+          {g_advance}
           beta_ += Hv;
-        }
-        for (int i = 0; i < n_per_t; ++i) {
+        }}
+        for (int i = 0; i < n_per_t; ++i) {{
           auto s_idx = n_per_t * dk_idx + i;
           o_state[s_idx] = static_cast<InT>(state[i]);
-        }
+        }}
     """
+    inputs = ["q", "k", "v", "g", "beta", "state_in", "T"]
+    if has_mask:
+        inputs.append("mask")
+
+    suffix = ""
+    if vectorized:
+        suffix += "_vec"
+    if has_mask:
+        suffix += "_mask"
+
     return mx.fast.metal_kernel(
-        name="gated_delta_step",
-        input_names=["q", "k", "v", "g", "beta", "state_in", "T"],
+        name=f"gated_delta_step{suffix}",
+        input_names=inputs,
         output_names=["y", "state_out"],
         source=source,
     )
 
 
-_gated_delta_kernel = _make_gated_delta_kernel()
+_gated_delta_kernel = _make_gated_delta_kernel(has_mask=False, vectorized=False)
+_gated_delta_kernel_masked = _make_gated_delta_kernel(has_mask=True, vectorized=False)
+_gated_delta_kernel_vec = _make_gated_delta_kernel(has_mask=False, vectorized=True)
+_gated_delta_kernel_vec_masked = _make_gated_delta_kernel(
+    has_mask=True, vectorized=True
+)
 
 
+@mx.compile
 def _gated_delta_step_ops(
     q: mx.array,
     k: mx.array,
@@ -99,6 +131,7 @@ def _gated_delta_step_ops(
     g: mx.array,
     beta: mx.array,
     state: mx.array,
+    mask: Optional[mx.array] = None,
 ) -> Tuple[mx.array, mx.array]:
     """
     Ops-based reference implementation for a single recurrent step.
@@ -106,7 +139,8 @@ def _gated_delta_step_ops(
     Shapes:
       - q, k: [B, H, Dk]
       - v: [B, H, Dv]
-      - g, beta: [B, H]
+      - g: [B, H] or [B, H, Dk]
+      - beta: [B, H]
       - state: [B, H, Dv, Dk]
     Returns:
       - y: [B, H, Dv]
@@ -114,12 +148,25 @@ def _gated_delta_step_ops(
     """
 
     # Decay
-    state = state * g[..., None, None]
+    old_state = state
+    if g.ndim == 2:
+        decay = g[..., None, None]
+    elif g.ndim == 3:
+        decay = g[..., None, :]
+    else:
+        raise ValueError(f"Unsupported gating shape {g.shape}")
+    state = state * decay
     kv_mem = (state * k[..., None, :]).sum(axis=-1)  # [B, H, Dv]
     delta = (v - kv_mem) * beta[..., None]  # [B, H, Dv]
     state = state + k[..., None, :] * delta[..., None]
     # Output projection along key dim with q
     y = (state * q[..., None, :]).sum(axis=-1)  # [B, H, Dv]
+    if mask is not None:
+        if mask.ndim == 2:
+            mask = mx.expand_dims(mask, axes=(2, 3))
+        elif mask.ndim == 3:
+            mask = mx.expand_dims(mask, axis=-1)
+        state = mx.where(mask, state, old_state)
     return y, state
 
 
@@ -130,12 +177,26 @@ def gated_delta_kernel(
     g: mx.array,
     beta: mx.array,
     state: mx.array,
+    mask: Optional[mx.array] = None,
 ) -> Tuple[mx.array, mx.array]:
     B, T, Hk, Dk = k.shape
     Hv, Dv = v.shape[2:]
     input_type = q.dtype
-    return _gated_delta_kernel(
-        inputs=[q, k, v, g, beta, state, T],
+    if g.ndim == 4:
+        kernel = _gated_delta_kernel_vec
+        inputs = [q, k, v, g, beta, state, T]
+        if mask is not None:
+            kernel = _gated_delta_kernel_vec_masked
+            inputs.append(mask)
+    else:
+        kernel = _gated_delta_kernel
+        inputs = [q, k, v, g, beta, state, T]
+        if mask is not None:
+            kernel = _gated_delta_kernel_masked
+            inputs.append(mask)
+
+    return kernel(
+        inputs=inputs,
         template=[
             ("InT", input_type),
             ("Dk", Dk),
@@ -157,18 +218,21 @@ def gated_delta_ops(
     g: mx.array,
     beta: mx.array,
     state: Optional[mx.array] = None,
+    mask: Optional[mx.array] = None,
 ) -> Tuple[mx.array, mx.array]:
     """
     Ops-based reference implementation for prompt prefill (sequential loop).
+    Supports both scalar and vectorized gating.
 
     Shapes:
       - q, k: [B, T, Hk, Dk]
       - v: [B, T, Hv, Dv]
-      - g, beta: [B, T, Hv]
-      - state: [B, Hv, Dk, Dv]
+      - g: [B, T, Hv] (scalar) or [B, T, Hv, Dk] (vectorized)
+      - beta: [B, T, Hv]
+      - state: [B, Hv, Dv, Dk]
     Returns:
       - y: [B, T, Hv, Dv]
-      - state: [B, Hv, Dk, Dv]
+      - state: [B, Hv, Dv, Dk]
     """
     B, T, Hk, Dk = q.shape
     Hv, Dv = v.shape[-2:]
@@ -188,6 +252,7 @@ def gated_delta_ops(
             g[:, t],
             beta[:, t],
             state,
+            None if mask is None else mask[:, t],
         )
         ys.append(y)
     y = mx.stack(ys, axis=1)
@@ -203,6 +268,8 @@ def gated_delta_update(
     A_log: mx.array,
     dt_bias: mx.array,
     state: Optional[mx.array] = None,
+    mask: Optional[mx.array] = None,
+    use_kernel: bool = True,
 ) -> Tuple[mx.array, mx.array]:
 
     beta = mx.sigmoid(b)
@@ -210,10 +277,8 @@ def gated_delta_update(
     if state is None:
         B, _, Hk, Dk = q.shape
         Hv, Dv = v.shape[-2:]
-        if state is None:
-            state = mx.zeros((B, Hv, Dv, Dk), dtype=q.dtype)
+        state = mx.zeros((B, Hv, Dv, Dk), dtype=q.dtype)
 
-    if mx.default_device() != mx.gpu or not mx.metal.is_available():
-        return gated_delta_ops(q, k, v, g, beta, state)
-    else:
-        return gated_delta_kernel(q, k, v, g, beta, state)
+    if not use_kernel or mx.default_device() != mx.gpu or not mx.metal.is_available():
+        return gated_delta_ops(q, k, v, g, beta, state, mask)
+    return gated_delta_kernel(q, k, v, g, beta, state, mask)

mlx_lm/models/gemma3_text.py

@@ -2,13 +2,14 @@
 
 from dataclasses import dataclass
 from functools import partial
-from typing import Any, Optional
+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 .cache import KVCache, RotatingKVCache
+from .rope_utils import initialize_rope
 
 
 @dataclass
@@ -22,12 +23,13 @@ class ModelArgs(BaseModelArgs):
     rms_norm_eps: float = 1.0e-6
     vocab_size: int = 262144
     num_key_value_heads: int = 1
-    rope_global_base_freq: float = 1_000_000.0
+    rope_theta: float = 1_000_000.0
     rope_local_base_freq: float = 10_000.0
-    rope_traditional: bool = False
     query_pre_attn_scalar: float = 256
     sliding_window: int = 512
     sliding_window_pattern: int = 6
+    max_position_embeddings: int = 32768
+    rope_scaling: Dict = None
 
 
 class Attention(nn.Module):
@@ -52,14 +54,12 @@ class Attention(nn.Module):
         self.k_norm = RMSNorm(dims=head_dim, eps=args.rms_norm_eps)
         self.is_sliding = (layer_idx + 1) % args.sliding_window_pattern != 0
 
-        self.rope = nn.RoPE(
-            head_dim,
-            traditional=args.rope_traditional,
-            base=(
-                args.rope_local_base_freq
-                if self.is_sliding
-                else args.rope_global_base_freq
-            ),
+        self.rope = initialize_rope(
+            dims=head_dim,
+            base=(args.rope_local_base_freq if self.is_sliding else args.rope_theta),
+            traditional=False,
+            max_position_embeddings=args.max_position_embeddings,
+            scaling_config=args.rope_scaling,
         )
 
     def __call__(
@@ -87,8 +87,6 @@ class Attention(nn.Module):
             keys = self.rope(keys)
 
         # Sliding window
-        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
         )
@@ -189,12 +187,14 @@ class Gemma3Model(nn.Module):
 
         global_mask = create_attention_mask(h, cache[self.sliding_window_pattern - 1])
 
-        sliding_window_mask = create_attention_mask(
-            h,
-            cache[0],
-            window_size=self.window_size,
-        )
-
+        if self.sliding_window_pattern > 1:
+            sliding_window_mask = create_attention_mask(
+                h,
+                cache[0],
+                window_size=self.window_size,
+            )
+        else:
+            sliding_window_mask = None
         for i, (layer, c) in enumerate(zip(self.layers, cache)):
             is_global = (
                 i % self.sliding_window_pattern == self.sliding_window_pattern - 1
@@ -246,7 +246,5 @@ class Model(nn.Module):
             ):
                 caches.append(KVCache())
             else:
-                caches.append(
-                    RotatingKVCache(max_size=self.args.sliding_window, keep=0)
-                )
+                caches.append(RotatingKVCache(max_size=self.args.sliding_window))
         return caches

mlx_lm/models/glm4_moe.py

@@ -9,6 +9,7 @@ import mlx.core as mx
 import mlx.nn as nn
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .pipeline import PipelineMixin
 from .switch_layers import SwitchGLU
 
 
@@ -243,7 +244,7 @@ class DecoderLayer(nn.Module):
         return h + r
 
 
-class LanguageModel(nn.Module):
+class LanguageModel(PipelineMixin, nn.Module):
     def __init__(self, config: ModelArgs):
         super().__init__()
         self.vocab_size = config.vocab_size
@@ -264,13 +265,28 @@ class LanguageModel(nn.Module):
     ) -> mx.array:
         h = self.embed_tokens(x)
 
-        if cache is None:
-            cache = [None] * self.num_layers
+        pipeline_rank = self.pipeline_rank
+        pipeline_size = self.pipeline_size
 
+        if cache is None:
+            cache = [None] * len(self.pipeline_layers)
         mask = create_attention_mask(h, cache[0])
 
-        for i in range(self.num_layers):
-            h = self.layers[self.start_idx + i](h, mask, cache[i])
+        # Receive from the previous process in the pipeline
+        if pipeline_rank < pipeline_size - 1:
+            h = mx.distributed.recv_like(h, (pipeline_rank + 1))
+
+        for l, c in zip(self.pipeline_layers, cache):
+            h = l(h, mask, cache=c)
+
+        # Send to the next process in the pipeline
+        if pipeline_rank != 0:
+            h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
+            if cache[-1] is not None:
+                cache[-1].keys = mx.depends(cache[-1].keys, h)
+
+        # Broadcast h while keeping it in the graph
+        h = mx.distributed.all_gather(h)[: h.shape[0]]
 
         return self.norm(h)
 
@@ -315,7 +331,7 @@ class Model(nn.Module):
 
     @property
     def layers(self):
-        return self.model.layers
+        return self.model.pipeline_layers
 
     @property
     def cast_predicate(self):

mlx_lm/models/gpt2.py

@@ -1,4 +1,4 @@
-# Copyright © 2023-2024 Apple Inc.
+# Copyright © 2023 - 2024 Apple Inc.
 
 from dataclasses import dataclass
 from typing import Any, Optional
@@ -138,7 +138,9 @@ class GPT2Model(nn.Module):
         if cache[0] is not None:
             offset = cache[0].offset
 
-        position_ids = mx.arange(offset, offset + L)
+        offset = mx.array(offset)
+        position_ids = mx.arange(L) + offset[..., None]
+
         hidden_states += self.wpe(position_ids)
 
         mask = create_attention_mask(hidden_states, cache[0])

mlx_lm/models/gpt_neox.py

@@ -23,6 +23,7 @@ class ModelArgs(BaseModelArgs):
     vocab_size: int
     rotary_emb_base: int
     rotary_pct: float
+    use_parallel_residual: bool = True
     num_key_value_heads: int = None
 
     def __post_init__(self):
@@ -107,6 +108,7 @@ class TransformerBlock(nn.Module):
         self.layer_norm_eps = args.layer_norm_eps
         self.attention = Attention(args)
         self.mlp = MLP(args)
+        self.use_parallel_residual = args.use_parallel_residual
         self.input_layernorm = nn.LayerNorm(
             self.hidden_size,
             eps=self.layer_norm_eps,
@@ -121,12 +123,20 @@ class TransformerBlock(nn.Module):
         mask: Optional[mx.array] = None,
         cache: Optional[Any] = None,
     ) -> mx.array:
-        residual = x
-        # NeoX runs attention and feedforward network in parallel.
-        attn = self.attention(self.input_layernorm(x), mask, cache)
-        ffn = self.mlp(self.post_attention_layernorm(x))
-        out = attn + ffn + residual
-        return out
+        if self.use_parallel_residual:
+            residual = x
+            # Run attention and feedforward network in parallel.
+            attn = self.attention(self.input_layernorm(x), mask, cache)
+            ffn = self.mlp(self.post_attention_layernorm(x))
+            out = attn + ffn + residual
+            return out
+        else:
+            # Run attention and feedforward network sequentially.
+            attn_output = self.attention(self.input_layernorm(x), mask, cache)
+            x = x + attn_output
+            ffn_output = self.mlp(self.post_attention_layernorm(x))
+            x = x + ffn_output
+            return x
 
 
 class GPTNeoXModel(nn.Module):

mlx_lm/models/granitemoehybrid.py

@@ -20,6 +20,7 @@ from .switch_layers import SwitchGLU
 
 @dataclass
 class ModelArgs(BaseModelArgs):
+    # Required fields (no defaults)
     model_type: str
     vocab_size: int
     hidden_size: int
@@ -29,34 +30,42 @@ class ModelArgs(BaseModelArgs):
     num_attention_heads: int
     num_key_value_heads: int
     attention_bias: bool
-
-    # Scalar multipliers
     embedding_multiplier: float
     attention_multiplier: float
     logits_scaling: float
     residual_multiplier: float
-
-    # MoE parameters
-    num_local_experts: int
-    num_experts_per_tok: int
-    shared_intermediate_size: int
-
-    # Mamba parameters
-    mamba_n_heads: int
-    mamba_d_head: int
-    mamba_proj_bias: bool
-    mamba_d_state: int
-    mamba_d_conv: int
-    mamba_n_groups: int
-    mamba_conv_bias: bool
-
     layer_types: List[str]
     rms_norm_eps: float
     rope_theta: float
+
+    # Optional fields (with defaults)
+    # MoE parameters (optional for dense mode)
+    num_local_experts: Optional[int] = None
+    num_experts_per_tok: Optional[int] = None
+    shared_intermediate_size: Optional[int] = None
+
+    # Mamba parameters (optional for non-hybrid mode)
+    mamba_n_heads: Optional[int] = None
+    mamba_d_head: Optional[int] = None
+    mamba_proj_bias: Optional[bool] = None
+    mamba_d_state: Optional[int] = None
+    mamba_d_conv: Optional[int] = None
+    mamba_n_groups: Optional[int] = None
+    mamba_conv_bias: Optional[bool] = None
+
+    # Dense MLP parameters (for non-MoE mode)
+    mlp_bias: bool = False
+
+    # Other optional parameters
     position_embedding_type: str = "rope"
     tie_word_embeddings: bool = True
     time_step_limit: Tuple[float, float] = (0.001, 100.0)
 
+    # Mode flags - inferred from num_local_experts
+    @property
+    def use_moe(self) -> bool:
+        return bool(self.num_local_experts)
+
 
 class GraniteMoeHybridRMSNormGated(nn.Module):
     def __init__(self, hidden_size: int, eps: float = 1e-6):
@@ -314,11 +323,27 @@ class GraniteMoeHybridSharedMLP(nn.Module):
         return self.output_linear(nn.silu(gate) * up)
 
 
+class GraniteMoeHybridMLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        dim = args.hidden_size
+        hidden_dim = args.intermediate_size
+        mlp_bias = args.mlp_bias
+
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=mlp_bias)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=mlp_bias)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=mlp_bias)
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
 class GraniteMoeHybridLayer(nn.Module):
     def __init__(self, args: ModelArgs, layer_type: str):
         super().__init__()
         self.layer_type = layer_type
         self.residual_multiplier = args.residual_multiplier
+        self.use_moe = args.use_moe
 
         self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
 
@@ -329,8 +354,14 @@ class GraniteMoeHybridLayer(nn.Module):
         else:
             raise ValueError(f"Unknown layer type: {layer_type}")
 
-        self.shared_mlp = GraniteMoeHybridSharedMLP(args)
-        self.block_sparse_moe = GraniteMoeHybridMoE(args)
+        # MoE or dense MLP after attention/mamba
+        if self.use_moe:
+            self.shared_mlp = GraniteMoeHybridSharedMLP(args)
+            self.block_sparse_moe = GraniteMoeHybridMoE(args)
+        else:
+            # Dense MLP mode
+            self.mlp = GraniteMoeHybridMLP(args)
+
         self.post_attention_layernorm = nn.RMSNorm(
             args.hidden_size, eps=args.rms_norm_eps
         )
@@ -352,13 +383,16 @@ class GraniteMoeHybridLayer(nn.Module):
 
         hidden_states = residual + hidden_states * self.residual_multiplier
 
-        # Second block: MoE + shared_mlp
+        # Second block: MoE + shared_mlp OR dense MLP
         residual = hidden_states
         normed = self.post_attention_layernorm(hidden_states)
 
-        moe_out = self.block_sparse_moe(normed)
-        shared_out = self.shared_mlp(normed)
-        mlp_out = moe_out + shared_out
+        if self.use_moe:
+            moe_out = self.block_sparse_moe(normed)
+            shared_out = self.shared_mlp(normed)
+            mlp_out = moe_out + shared_out
+        else:
+            mlp_out = self.mlp(normed)
 
         hidden_states = residual + mlp_out * self.residual_multiplier
 
@@ -375,9 +409,16 @@ class GraniteMoeHybridModel(nn.Module):
         ]
         self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
         self.embedding_multiplier = args.embedding_multiplier
-        self.fa_idx = args.layer_types.index("attention")
-        self.ssm_idx = args.layer_types.index("mamba")
-        self.layer_types = args.layer_types
+
+        # Handle hybrid vs non-hybrid mode
+        self.fa_idx = (
+            args.layer_types.index("attention")
+            if "attention" in args.layer_types
+            else None
+        )
+        self.ssm_idx = (
+            args.layer_types.index("mamba") if "mamba" in args.layer_types else None
+        )
 
     def __call__(
         self,
@@ -389,11 +430,16 @@ class GraniteMoeHybridModel(nn.Module):
         if cache is None:
             cache = [None] * len(self.layers)
 
-        attn_mask = create_attention_mask(hidden_states, cache[self.fa_idx])
-        mamba_mask = create_ssm_mask(hidden_states, cache[self.ssm_idx])
+        # Create masks based on what layer types exist
+        attn_mask = None
+        mamba_mask = None
 
-        cache_counter = 0
-        for layer, c, layer_type in zip(self.layers, cache, self.layer_types):
+        if self.fa_idx is not None:
+            attn_mask = create_attention_mask(hidden_states, cache[self.fa_idx])
+        if self.ssm_idx is not None:
+            mamba_mask = create_ssm_mask(hidden_states, cache[self.ssm_idx])
+
+        for layer, c in zip(self.layers, cache):
             mask = attn_mask if layer.layer_type == "attention" else mamba_mask
             hidden_states = layer(hidden_states, mask=mask, cache=c)
 
@@ -443,8 +489,11 @@ class Model(nn.Module):
             if "conv1d.weight" in k and v.shape[-1] != 1:
                 weights[k] = v.moveaxis(2, 1)
 
-        # Handle MoE weight transformation to SwitchGLU format
-        if "model.layers.0.block_sparse_moe.input_linear.weight" in weights:
+        # Handle MoE weight transformation to SwitchGLU format (only for MoE models)
+        if (
+            self.args.use_moe
+            and "model.layers.0.block_sparse_moe.input_linear.weight" in weights
+        ):
             for l in range(self.args.num_hidden_layers):
                 prefix = f"model.layers.{l}.block_sparse_moe"
 
@@ -461,12 +510,31 @@ class Model(nn.Module):
                     f"{prefix}.output_linear.weight"
                 )
 
+        # Handle dense MLP weight transformation (for dense models)
+        elif (
+            not self.args.use_moe
+            and "model.layers.0.shared_mlp.input_linear.weight" in weights
+        ):
+            for l in range(self.args.num_hidden_layers):
+                prefix = f"model.layers.{l}.shared_mlp"
+
+                # Transform shared_mlp weights to standard mlp weights
+                input_weight = weights.pop(f"{prefix}.input_linear.weight")
+                # Split into gate and up projections (each half)
+                gate_proj, up_proj = mx.split(input_weight, 2, axis=0)
+                weights[f"model.layers.{l}.mlp.gate_proj.weight"] = gate_proj
+                weights[f"model.layers.{l}.mlp.up_proj.weight"] = up_proj
+
+                weights[f"model.layers.{l}.mlp.down_proj.weight"] = weights.pop(
+                    f"{prefix}.output_linear.weight"
+                )
+
         return weights
 
     @property
     def quant_predicate(self):
         def predicate(path, _):
-            if path.endswith("router.layer"):
+            if self.args.use_moe and path.endswith("router.layer"):
                 return {"group_size": 64, "bits": 8}
             return True
 

mlx_lm/models/jamba.py

@@ -0,0 +1,384 @@
+# Copyright © 2025 Apple Inc.
+
+import math
+from dataclasses import dataclass
+from typing import Any, List, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import (
+    BaseModelArgs,
+    create_attention_mask,
+    create_ssm_mask,
+    scaled_dot_product_attention,
+)
+from .cache import KVCache, MambaCache
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    attn_layer_offset: int
+    attn_layer_period: int
+    expert_layer_offset: int
+    expert_layer_period: int
+    mamba_d_conv: int
+    mamba_d_state: int
+    mamba_expand: int
+    num_experts: int
+    num_experts_per_tok: int
+    rms_norm_eps: float
+    max_position_embeddings: int
+    vocab_size: int
+    mamba_dt_rank: Union[str, int] = "auto"
+    mamba_proj_bias: bool = False
+    mamba_conv_bias: bool = True
+    layers_block_type: Optional[List[str]] = None
+    tie_word_embeddings: bool = True
+
+    def __post_init__(self):
+        if self.mamba_dt_rank == "auto":
+            self.mamba_dt_rank = math.ceil(self.hidden_size / 16)
+        if self.layers_block_type is None:
+            self.layers_block_type = [
+                (
+                    "attention"
+                    if i % self.attn_layer_period == self.attn_layer_offset
+                    else "mamba"
+                )
+                for i in range(self.num_hidden_layers)
+            ]
+
+
+class JambaMLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.gate_proj = nn.Linear(args.hidden_size, args.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(args.hidden_size, args.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(args.intermediate_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 JambaAttention(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
+        )
+
+    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.num_attention_heads, -1).transpose(
+            0, 2, 1, 3
+        )
+        keys = 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:
+            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)
+
+
+@mx.compile
+def fma(a, b, c):
+    return a * b + c
+
+
+class JambaMambaMixer(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.hidden_size = args.hidden_size
+        self.ssm_state_size = args.mamba_d_state
+        self.conv_kernel_size = args.mamba_d_conv
+        self.intermediate_size = args.mamba_expand * args.hidden_size
+        self.time_step_rank = args.mamba_dt_rank
+        self.use_conv_bias = args.mamba_conv_bias
+        self.use_bias = args.mamba_proj_bias
+
+        self.in_proj = nn.Linear(
+            self.hidden_size, self.intermediate_size * 2, bias=self.use_bias
+        )
+
+        self.conv1d = nn.Conv1d(
+            in_channels=self.intermediate_size,
+            out_channels=self.intermediate_size,
+            kernel_size=self.conv_kernel_size,
+            groups=self.intermediate_size,
+            bias=self.use_conv_bias,
+            padding=0,
+        )
+        self.x_proj = nn.Linear(
+            self.intermediate_size,
+            self.time_step_rank + self.ssm_state_size * 2,
+            bias=False,
+        )
+        self.dt_proj = nn.Linear(self.time_step_rank, self.intermediate_size, bias=True)
+
+        A = mx.repeat(
+            mx.arange(1.0, self.ssm_state_size + 1.0).reshape([1, self.ssm_state_size]),
+            repeats=self.intermediate_size,
+            axis=0,
+        )
+        self.A_log = mx.log(A)
+        self.D = mx.ones([self.intermediate_size])
+
+        self.out_proj = nn.Linear(
+            self.intermediate_size, self.hidden_size, bias=self.use_bias
+        )
+
+        self.dt_layernorm = nn.RMSNorm(self.time_step_rank, eps=args.rms_norm_eps)
+        self.b_layernorm = nn.RMSNorm(self.ssm_state_size, eps=args.rms_norm_eps)
+        self.c_layernorm = nn.RMSNorm(self.ssm_state_size, eps=args.rms_norm_eps)
+
+    def ssm_step(self, x, A, state=None):
+        T = x.shape[1]
+        D = self.D
+        deltaBC = self.x_proj(x)
+        delta, B, C = mx.split(
+            deltaBC,
+            [self.time_step_rank, self.time_step_rank + self.ssm_state_size],
+            axis=-1,
+        )
+        delta, B, C = self.dt_layernorm(delta), self.b_layernorm(B), self.c_layernorm(C)
+        delta = nn.softplus(self.dt_proj(delta))
+        new_state = mx.expand_dims(delta * x, -1) * mx.expand_dims(B, -2)
+        dtA = mx.exp(mx.expand_dims(delta, -1) * A)
+
+        # TODO, speed up prefill with chunked scan
+        for t in range(T):
+            if state is not None:
+                new_state[:, t] = fma(state, dtA[:, t], new_state[:, t])
+            state = new_state[:, t]
+        y = (new_state @ mx.expand_dims(C, -1)).squeeze(-1)
+        y = y + D * x
+        return y, new_state[:, -1]
+
+    def _process_sequence(self, x, conv_state, ssm_state):
+        xz = self.in_proj(x)
+        x, z = xz.split(indices_or_sections=2, axis=-1)
+        K = self.conv_kernel_size
+        if conv_state is not None:
+            x_full = mx.concatenate([conv_state, x], axis=1)
+        else:
+            x_full = mx.pad(x, [(0, 0), (K - 1, 0), (0, 0)])
+        conv_out = self.conv1d(x_full)
+        conv_state = x_full[:, -(K - 1) :, :]
+        x = nn.silu(conv_out)
+        A = -mx.exp(self.A_log)
+        y, ssm_state = self.ssm_step(x, A, ssm_state)
+        z = self.out_proj(nn.silu(z) * y)
+        return z, (conv_state, ssm_state)
+
+    def __call__(self, x, cache):
+        if cache is None:
+            conv_state, ssm_state = None, None
+        else:
+            conv_state, ssm_state = cache[0], cache[1]
+
+        output, (conv_state, ssm_state) = self._process_sequence(
+            x, conv_state, ssm_state
+        )
+
+        if cache is not None:
+            cache[0] = conv_state
+            cache[1] = ssm_state
+
+        return output
+
+
+class JambaSparseMoeBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_experts_per_tok = args.num_experts_per_tok
+
+        self.router = nn.Linear(args.hidden_size, args.num_experts, bias=False)
+        self.switch_mlp = SwitchGLU(
+            args.hidden_size, args.intermediate_size, args.num_experts
+        )
+
+    def __call__(self, x: mx.array) -> mx.array:
+        gates = self.router(x)
+        k = self.num_experts_per_tok
+        inds = mx.stop_gradient(mx.argpartition(-gates, kth=k - 1, axis=-1)[..., :k])
+        scores = mx.take_along_axis(gates, inds, axis=-1)
+        scores = mx.softmax(scores, axis=-1, precise=True)
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2)
+        return y
+
+
+class JambaDecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_type: str, layer_idx: int):
+        super().__init__()
+        self.is_attn = layer_type == "attention"
+        if self.is_attn:
+            self.self_attn = JambaAttention(args)
+        else:
+            self.mamba = JambaMambaMixer(args)
+        if (
+            args.num_experts > 1
+            and (layer_idx + args.expert_layer_offset) % args.expert_layer_period == 0
+        ):
+            ffn_layer_class = JambaSparseMoeBlock
+        else:
+            ffn_layer_class = JambaMLP
+        self.feed_forward = ffn_layer_class(args)
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.pre_ff_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:
+        if self.is_attn:
+            h = self.self_attn(self.input_layernorm(x), mask, cache)
+        else:
+            h = self.mamba(self.input_layernorm(x), cache)
+        r = x + h
+        out = r + self.feed_forward(self.pre_ff_layernorm(r))
+        return out
+
+
+class JambaModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+
+        self.layers = [
+            JambaDecoderLayer(args, t, idx)
+            for idx, t in enumerate(args.layers_block_type)
+        ]
+        self.final_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.attn_idx = args.layers_block_type.index("attention")
+        self.ssm_idx = args.layers_block_type.index("mamba")
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        h = self.embed_tokens(inputs)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        attn_mask = create_attention_mask(h, cache[self.attn_idx])
+        ssm_mask = create_ssm_mask(h, cache[self.ssm_idx])
+
+        for layer, c in zip(self.layers, cache):
+            mask = attn_mask if layer.is_attn else ssm_mask
+            h = layer(h, mask=mask, cache=c)
+
+        return self.final_layernorm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.model_type = args.model_type
+        self.args = args
+        self.model = JambaModel(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        out = self.model(inputs, cache)
+        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):
+        caches = []
+        for layer in self.model.layers:
+            if layer.is_attn:
+                caches.append(KVCache())
+            else:
+                caches.append(MambaCache())
+        return caches
+
+    def sanitize(self, weights):
+        for k, v in list(weights.items()):
+            if "conv1d.weight" in k and v.shape[-1] != 1:
+                weights[k] = v.moveaxis(2, 1)
+
+        if self.args.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
+
+        for l in range(self.args.num_hidden_layers):
+            base = f"model.layers.{l}.feed_forward"
+            if not any(key.startswith(f"{base}.experts.") for key in weights.keys()):
+                continue
+
+            for proj in ["gate_proj", "down_proj", "up_proj"]:
+                for name in ["weight", "bias", "scales", "biases"]:
+                    expert_tensors = [
+                        weights.pop(f"{base}.experts.{e}.{proj}.{name}")
+                        for e in range(len(weights))
+                        if f"{base}.experts.{e}.{proj}.{name}" in weights
+                    ]
+                    if expert_tensors:
+                        weights[f"{base}.switch_mlp.{proj}.{name}"] = mx.stack(
+                            expert_tensors
+                        )
+
+        return 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

mlx_lm/models/kimi_linear.py

@@ -0,0 +1,575 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import (
+    BaseModelArgs,
+    create_attention_mask,
+    create_ssm_mask,
+    scaled_dot_product_attention,
+)
+from .cache import KVCache, MambaCache
+from .gated_delta import gated_delta_update
+from .rope_utils import initialize_rope
+from .switch_layers import SwitchGLU
+
+
+@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
+    intermediate_size: int
+    head_dim: int
+    rope_theta: float
+    rms_norm_eps: float
+    linear_attn_config: Dict[str, Any]
+    model_max_length: int
+    num_experts: int
+    moe_intermediate_size: int
+    kv_lora_rank: int
+    rope_scaling: Optional[Dict[str, Any]] = None
+    tie_word_embeddings: bool = False
+    qk_nope_head_dim: Optional[int] = None
+    qk_rope_head_dim: Optional[int] = None
+    v_head_dim: Optional[int] = None
+    mla_use_nope: bool = False
+    num_experts_per_token: int = 1
+    num_shared_experts: int = 0
+    moe_router_activation_func: str = "sigmoid"
+    moe_renormalize: bool = True
+    routed_scaling_factor: float = 1.0
+    first_k_dense_replace: int = 0
+    moe_layer_freq: int = 1
+    use_grouped_topk: bool = True
+    num_expert_group: int = 1
+    topk_group: int = 1
+
+
+class KimiMLP(nn.Module):
+    def __init__(
+        self,
+        args: ModelArgs,
+        hidden_size: Optional[int] = None,
+        intermediate_size: Optional[int] = None,
+    ):
+        super().__init__()
+        dim = hidden_size or args.hidden_size
+        hidden = intermediate_size or args.intermediate_size
+        self.gate_proj = nn.Linear(dim, hidden, bias=False)
+        self.up_proj = nn.Linear(dim, hidden, bias=False)
+        self.down_proj = nn.Linear(hidden, dim, bias=False)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+@mx.compile
+def _group_expert_select(
+    gates: mx.array,
+    bias: Optional[mx.array],
+    top_k: int,
+    n_group: int,
+    topk_group: int,
+    routed_scaling_factor: float,
+    renormalize: bool,
+    score_function: str,
+) -> Tuple[mx.array, mx.array]:
+    if score_function == "sigmoid":
+        scores = mx.sigmoid(gates)
+    elif score_function == "softmax":
+        scores = mx.softmax(gates, axis=-1, precise=True)
+    else:
+        raise ValueError(f"Unsupported MoE router activation '{score_function}'")
+
+    orig_scores = scores
+    if bias is not None:
+        scores = scores + bias.astype(scores.dtype)
+
+    if n_group > 1:
+        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, dtype=scores.dtype),
+            axis=-2,
+        )
+        scores = mx.flatten(scores, -2, -1)
+
+    inds = mx.argpartition(-scores, kth=top_k - 1, axis=-1)[..., :top_k]
+    scores = mx.take_along_axis(orig_scores, inds, axis=-1)
+
+    if top_k > 1 and renormalize:
+        denominator = scores.sum(axis=-1, keepdims=True) + 1e-20
+        scores = scores / denominator
+
+    return inds, scores * routed_scaling_factor
+
+
+class KimiSparseMoE(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        hidden = args.hidden_size
+        experts = args.num_experts
+        if experts is None:
+            raise ValueError("num_experts must be specified for MoE layers")
+
+        self.gate = nn.Linear(hidden, experts, bias=False)
+        self.switch_mlp = SwitchGLU(hidden, args.moe_intermediate_size, experts)
+        self.e_score_correction_bias = mx.zeros((experts,), dtype=mx.float32)
+
+        if args.num_shared_experts:
+            shared_hidden = args.moe_intermediate_size * args.num_shared_experts
+            self.shared_experts = KimiMLP(args, intermediate_size=shared_hidden)
+        else:
+            self.shared_experts = None
+
+    def __call__(self, x: mx.array) -> mx.array:
+        scores = self.gate(x)
+        inds, weights = _group_expert_select(
+            scores,
+            self.e_score_correction_bias,
+            self.args.num_experts_per_token,
+            self.args.num_expert_group,
+            self.args.topk_group,
+            self.args.routed_scaling_factor,
+            self.args.moe_renormalize,
+            self.args.moe_router_activation_func,
+        )
+        out = self.switch_mlp(x, inds)
+        out = (out * weights[..., None]).sum(axis=-2)
+        if self.shared_experts is not None:
+            out = out + self.shared_experts(x)
+        return out
+
+
+class KimiMLAAttention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.num_heads = args.num_attention_heads
+        self.num_key_value_heads = args.num_key_value_heads
+        self.qk_nope_head_dim = args.qk_nope_head_dim or args.head_dim
+        self.qk_rope_head_dim = args.qk_rope_head_dim or 0
+        self.q_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+        self.v_head_dim = args.v_head_dim or args.head_dim
+        self.scale = self.q_head_dim**-0.5
+
+        hidden = args.hidden_size
+        self.q_proj = nn.Linear(hidden, self.num_heads * self.q_head_dim, bias=False)
+        self.kv_a_proj_with_mqa = nn.Linear(
+            hidden,
+            args.kv_lora_rank + self.qk_rope_head_dim,
+            bias=False,
+        )
+        self.kv_a_layernorm = nn.RMSNorm(args.kv_lora_rank, eps=args.rms_norm_eps)
+        self.kv_b_proj = nn.Linear(
+            args.kv_lora_rank,
+            self.num_heads
+            * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
+            bias=False,
+        )
+        self.o_proj = nn.Linear(self.num_heads * self.v_head_dim, hidden, bias=False)
+
+        rope_dim = self.qk_rope_head_dim or self.q_head_dim
+        self.rope = initialize_rope(
+            rope_dim,
+            base=args.rope_theta,
+            traditional=False,
+            scaling_config=args.rope_scaling,
+            max_position_embeddings=args.model_max_length,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[KVCache] = None,
+    ) -> mx.array:
+        B, L, _ = x.shape
+        q_states = self.q_proj(x).reshape(B, L, self.num_heads, self.q_head_dim)
+        q_pass, q_rot = mx.split(q_states, [self.qk_nope_head_dim], axis=-1)
+
+        compressed = self.kv_a_proj_with_mqa(x)
+        k_pass, k_rot = mx.split(
+            compressed, [compressed.shape[-1] - self.qk_rope_head_dim], axis=-1
+        )
+        k_pass = self.kv_a_layernorm(k_pass)
+        kv = self.kv_b_proj(k_pass)
+        kv = kv.reshape(
+            B,
+            L,
+            self.num_heads,
+            self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim,
+        )
+        k_pass, v_states = mx.split(kv, [self.qk_nope_head_dim], axis=-1)
+
+        if self.qk_rope_head_dim:
+            k_rot = mx.reshape(k_rot, (B, L, 1, self.qk_rope_head_dim))
+            k_rot = mx.broadcast_to(k_rot, (*k_pass.shape[:-1], self.qk_rope_head_dim))
+        else:
+            k_rot = mx.zeros((*k_pass.shape[:-1], 0), dtype=k_pass.dtype)
+
+        queries = mx.concatenate([q_pass, q_rot], axis=-1).transpose(0, 2, 1, 3)
+        keys = mx.concatenate([k_pass, k_rot], axis=-1).transpose(0, 2, 1, 3)
+        values = v_states.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)
+
+        out = scaled_dot_product_attention(
+            queries,
+            keys,
+            values,
+            cache,
+            scale=self.scale,
+            mask=mask,
+        )
+        out = out.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(out)
+
+
+class ShortConv1d(nn.Module):
+    def __init__(self, channels: int, kernel_size: int):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.conv = nn.Conv1d(
+            in_channels=channels,
+            out_channels=channels,
+            kernel_size=kernel_size,
+            bias=False,
+            groups=channels,
+            padding=0,
+        )
+
+    def __call__(
+        self, x: mx.array, cache: Optional[mx.array]
+    ) -> Tuple[mx.array, mx.array]:
+        if cache is None:
+            pad = mx.zeros(
+                (x.shape[0], self.kernel_size - 1, x.shape[-1]), dtype=x.dtype
+            )
+        else:
+            pad = cache
+        conv_input = mx.concatenate([pad, x], axis=1)
+        out = nn.silu(self.conv(conv_input))
+        new_cache = conv_input[:, -self.kernel_size + 1 :, :]
+        return out, new_cache
+
+
+class KimiDeltaAttention(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        cfg = args.linear_attn_config
+
+        self.layer_idx = layer_idx
+        self.num_heads = cfg["num_heads"]
+        self.head_dim = cfg["head_dim"]
+        self.conv_kernel = cfg.get("short_conv_kernel_size", 4)
+
+        self.projection_dim = self.num_heads * self.head_dim
+        hidden = args.hidden_size
+
+        self.scale = float(self.head_dim) ** -0.5
+
+        self.q_proj = nn.Linear(hidden, self.projection_dim, bias=False)
+        self.k_proj = nn.Linear(hidden, self.projection_dim, bias=False)
+        self.v_proj = nn.Linear(hidden, self.projection_dim, bias=False)
+
+        self.q_conv = ShortConv1d(self.projection_dim, self.conv_kernel)
+        self.k_conv = ShortConv1d(self.projection_dim, self.conv_kernel)
+        self.v_conv = ShortConv1d(self.projection_dim, self.conv_kernel)
+
+        self.f_a_proj = nn.Linear(hidden, self.head_dim, bias=False)
+        self.f_b_proj = nn.Linear(self.head_dim, self.projection_dim, bias=False)
+        self.b_proj = nn.Linear(hidden, self.num_heads, bias=False)
+
+        self.g_a_proj = nn.Linear(hidden, self.head_dim, bias=False)
+        self.g_b_proj = nn.Linear(self.head_dim, self.projection_dim, bias=False)
+
+        self.A_log = mx.expand_dims(
+            mx.log(mx.random.uniform(low=1.0, high=16.0, shape=(self.num_heads,))),
+            (0, 1, 3),
+        )
+        self.dt_bias = mx.zeros((self.projection_dim,))
+
+        self.o_norm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
+        self.o_proj = nn.Linear(self.projection_dim, hidden, bias=False)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, T, _ = x.shape
+        dtype = x.dtype
+
+        if cache is not None:
+            conv_state, ssm_state = cache
+        else:
+            conv_state = None
+            ssm_state = None
+
+        if conv_state is None:
+            s = mx.zeros((B, self.conv_kernel - 1, self.projection_dim), dtype=dtype)
+            q_state = s
+            k_state = s
+            v_state = s
+        else:
+            q_state, k_state, v_state = conv_state
+
+        q_conv, q_state = self.q_conv(self.q_proj(x), q_state)
+        k_conv, k_state = self.k_conv(self.k_proj(x), k_state)
+        v_conv, v_state = self.v_conv(self.v_proj(x), v_state)
+
+        if cache is not None:
+            cache[0] = (q_state, k_state, v_state)
+
+        q = q_conv.reshape(B, T, self.num_heads, self.head_dim)
+        k = k_conv.reshape(B, T, self.num_heads, self.head_dim)
+        v = v_conv.reshape(B, T, self.num_heads, self.head_dim)
+
+        def _l2norm(x, eps=1e-6):
+            norm = mx.linalg.norm(x, axis=-1, keepdims=True)
+            return x / (norm + eps)
+
+        q = _l2norm(q)
+        k = _l2norm(k)
+        q = q * self.scale
+
+        a_logits = self.f_b_proj(self.f_a_proj(x)).reshape(
+            B, T, self.num_heads, self.head_dim
+        )
+        b_logits = self.b_proj(x).reshape(B, T, self.num_heads)
+
+        out, ssm_state = gated_delta_update(
+            q,
+            k,
+            v,
+            a_logits,
+            b_logits,
+            self.A_log.reshape(self.num_heads, 1),
+            self.dt_bias.reshape(self.num_heads, self.head_dim),
+            state=ssm_state,
+            mask=mask,
+            use_kernel=not self.training,
+        )
+
+        if cache is not None:
+            cache[1] = ssm_state
+
+        gate = self.g_b_proj(self.g_a_proj(x)).reshape(
+            B, T, self.num_heads, self.head_dim
+        )
+        out = (
+            self.o_norm(out.reshape(B, T, self.num_heads, self.head_dim))
+            * mx.sigmoid(gate)
+        ).reshape(B, T, -1)
+        return self.o_proj(out)
+
+
+class KimiDecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        kda_layers = args.linear_attn_config["kda_layers"]
+        self.is_linear = (layer_idx + 1) in kda_layers
+
+        if self.is_linear:
+            self.self_attn = KimiDeltaAttention(args, layer_idx)
+        else:
+            self.self_attn = KimiMLAAttention(args)
+
+        if (
+            args.num_experts > 0
+            and layer_idx >= args.first_k_dense_replace
+            and layer_idx % args.moe_layer_freq == 0
+        ):
+            self.mlp = KimiSparseMoE(args)
+        else:
+            self.mlp = KimiMLP(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:
+        attn_cache = None if cache is None else cache
+        y = self.self_attn(self.input_layernorm(x), mask, attn_cache)
+        h = x + y
+        z = self.mlp(self.post_attention_layernorm(h))
+        return h + z
+
+
+class KimiLinearModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [KimiDecoderLayer(args, i) for i in range(args.num_hidden_layers)]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        kda_layers = args.linear_attn_config["kda_layers"]
+        self.ssm_idx = kda_layers[0] - 1
+        for i in range(len(self.layers)):
+            if (i + 1) not in kda_layers:
+                self.attn_idx = i
+                break
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[List[Any]] = None,
+    ) -> mx.array:
+        h = self.embed_tokens(inputs)
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        ssm_mask = create_ssm_mask(h, cache[self.ssm_idx])
+        attn_mask = create_attention_mask(h, cache[self.attn_idx])
+
+        for layer, layer_cache in zip(self.layers, cache):
+            mask = ssm_mask if layer.is_linear else attn_mask
+            h = layer(h, mask=mask, cache=layer_cache)
+
+        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 = KimiLinearModel(args)
+        if args.tie_word_embeddings:
+            self.lm_head = None
+        else:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[List[Any]] = None,
+    ) -> mx.array:
+        out = self.model(inputs, cache)
+        if self.lm_head is None:
+            return self.model.embed_tokens.as_linear(out)
+        return self.lm_head(out)
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def make_cache(self):
+        caches: List[Any] = []
+        for layer in self.layers:
+            if layer.is_linear:
+                caches.append(MambaCache())
+            else:
+                caches.append(KVCache())
+        return caches
+
+    def sanitize(self, weights: Dict[str, mx.array]) -> Dict[str, mx.array]:
+        weights = {k: v for k, v in weights.items() if not k.startswith("model.mtp")}
+
+        if self.args.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
+
+        for layer_idx, layer in enumerate(self.layers):
+            prefix = f"model.layers.{layer_idx}"
+
+            if isinstance(layer.mlp, KimiSparseMoE):
+                src_prefix = f"{prefix}.block_sparse_moe"
+                dst_prefix = f"{prefix}.mlp"
+                for src, dst in [
+                    ("w1", "gate_proj"),
+                    ("w2", "down_proj"),
+                    ("w3", "up_proj"),
+                ]:
+                    key = f"{src_prefix}.experts.0.{src}.weight"
+                    if key in weights:
+                        stacked = [
+                            weights.pop(f"{src_prefix}.experts.{i}.{src}.weight")
+                            for i in range(self.args.num_experts)
+                        ]
+                        weights[f"{dst_prefix}.switch_mlp.{dst}.weight"] = mx.stack(
+                            stacked
+                        )
+
+                for name in ("gate_proj", "up_proj", "down_proj"):
+                    src_key = f"{src_prefix}.shared_experts.{name}.weight"
+                    if src_key in weights:
+                        weights[f"{dst_prefix}.shared_experts.{name}.weight"] = (
+                            weights.pop(src_key)
+                        )
+
+                gate_key = f"{src_prefix}.gate.weight"
+                if gate_key in weights:
+                    weights[f"{dst_prefix}.gate.weight"] = weights.pop(gate_key)
+
+                bias_key = f"{src_prefix}.gate.e_score_correction_bias"
+                if bias_key in weights:
+                    weights[f"{dst_prefix}.e_score_correction_bias"] = weights.pop(
+                        bias_key
+                    )
+
+            attn = getattr(layer, "self_attn", None)
+            if isinstance(attn, KimiDeltaAttention):
+                attn_prefix = f"{prefix}.self_attn"
+                for src_name, dst_name in (
+                    ("q_conv1d", "q_conv"),
+                    ("k_conv1d", "k_conv"),
+                    ("v_conv1d", "v_conv"),
+                ):
+                    src_key = f"{attn_prefix}.{src_name}.weight"
+                    if src_key in weights:
+                        w = weights.pop(src_key)
+                        if w.ndim == 3:
+                            w = w.moveaxis(2, 1)
+                        weights[f"{attn_prefix}.{dst_name}.conv.weight"] = w
+                dt_key = f"{attn_prefix}.dt_bias"
+                if dt_key in weights:
+                    if weights[dt_key].ndim > 1:
+                        weights[dt_key] = mx.reshape(weights[dt_key], (-1,))
+
+        return weights
+
+    @property
+    def cast_predicate(self):
+        def predicate(path: str):
+            if "e_score_correction_bias" in path:
+                return False
+            if path.endswith("A_log") or path.endswith("dt_bias"):
+                return False
+            return True
+
+        return predicate
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if path.endswith("mlp.gate"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate

mlx_lm/models/lfm2-vl.py

@@ -18,11 +18,6 @@ class ModelArgs(BaseModelArgs):
 
     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):

mlx_lm/models/lfm2.py

@@ -31,8 +31,19 @@ class ModelArgs(BaseModelArgs):
     block_multiple_of: int
     block_ffn_dim_multiplier: float
     block_auto_adjust_ff_dim: bool
-    full_attn_idxs: List[int]
     rope_theta: float
+    full_attn_idxs: Optional[List[int]] = None
+    layer_types: Optional[List[str]] = None
+
+    def __post_init__(self):
+        if self.num_key_value_heads is None:
+            self.num_key_value_heads = self.num_attention_heads
+        if self.full_attn_idxs is None:
+            self.full_attn_idxs = [
+                i
+                for i, layer_type in enumerate(self.layer_types)
+                if layer_type == "full_attention"
+            ]
 
 
 class Attention(nn.Module):
@@ -122,7 +133,6 @@ class ShortConv(nn.Module):
         mask: Optional[mx.array] = None,
         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

mlx_lm/models/lfm2_moe.py

@@ -0,0 +1,372 @@
+# 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,
+    create_ssm_mask,
+    scaled_dot_product_attention,
+)
+from .cache import ArraysCache, KVCache
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    moe_intermediate_size: int
+    num_hidden_layers: int
+    num_experts: int
+    num_experts_per_tok: int
+    norm_topk_prob: bool
+    num_attention_heads: int
+    num_key_value_heads: int
+    max_position_embeddings: int
+    use_expert_bias: bool
+    num_dense_layers: int
+    norm_eps: float
+    conv_bias: bool
+    conv_L_cache: int
+    rope_theta: float
+    full_attn_idxs: Optional[List[int]] = None
+    layer_types: Optional[List[str]] = None
+
+    def __post_init__(self):
+        if self.full_attn_idxs is None:
+            self.full_attn_idxs = [
+                i
+                for i, layer_type in enumerate(self.layer_types)
+                if layer_type == "full_attention"
+            ]
+
+
+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,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ):
+        BCx = self.in_proj(x)
+        B, C, x = mx.split(BCx, 3, axis=-1)
+        Bx = B * x
+        if mask is not None:
+            Bx = mx.where(mask[..., None], Bx, 0)
+        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, config: ModelArgs, intermediate_size: Optional[int] = None):
+        super().__init__()
+        self.hidden_size = config.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) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class Lfm2MoeSparseMoeBlock(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.use_expert_bias = args.use_expert_bias
+
+        self.gate = nn.Linear(dim, num_experts, bias=False)
+        self.switch_mlp = SwitchGLU(dim, intermediate_size, num_experts)
+        if self.use_expert_bias:
+            self.expert_bias = mx.zeros((self.num_experts,))
+
+    def __call__(
+        self,
+        x: mx.array,
+    ):
+        gates = self.gate(x).astype(mx.float32)
+        gates = mx.softmax(gates, axis=-1)
+
+        if self.use_expert_bias:
+            gates += self.expert_bias
+
+        k = self.top_k
+        inds = mx.argpartition(gates, kth=-k, 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) + 1e-20
+        scores = scores.astype(x.dtype)
+
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2)
+
+        return y
+
+
+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(
+                config=args,
+                intermediate_size=args.intermediate_size,
+            )
+            if layer_idx < args.num_dense_layers
+            else Lfm2MoeSparseMoeBlock(args)
+        )
+
+        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),
+                mask=mask,
+                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)
+
+        self.fa_idx = args.full_attn_idxs[0]
+        self.conv_idx = 0
+        for i in range(args.num_hidden_layers):
+            if i in args.full_attn_idxs:
+                self.conv_idx += 1
+            else:
+                break
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        if input_embeddings is not None:
+            h = input_embeddings
+        else:
+            h = self.embed_tokens(inputs)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        attn_mask = create_attention_mask(h, cache[self.fa_idx])
+        conv_mask = create_ssm_mask(h, cache[self.conv_idx])
+
+        for layer, c in zip(self.layers, cache):
+            mask = attn_mask if layer.is_attention_layer else conv_mask
+            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,
+        cache=None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        out = self.model(inputs, 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)
+            replacements = {
+                "w1.weight": "gate_proj.weight",
+                "w2.weight": "down_proj.weight",
+                "w3.weight": "up_proj.weight",
+            }
+            for old, new in replacements.items():
+                if old in name:
+                    name = name.replace(old, new)
+            sanitized_weights[name] = param
+
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}"
+            # Only sanitize MoE layer weights
+            for n in ["gate_proj", "down_proj", "up_proj"]:
+                if f"{prefix}.feed_forward.experts.0.{n}.weight" in sanitized_weights:
+                    to_join = [
+                        sanitized_weights.pop(
+                            f"{prefix}.feed_forward.experts.{e}.{n}.weight"
+                        )
+                        for e in range(self.args.num_experts)
+                    ]
+                    sanitized_weights[
+                        f"{prefix}.feed_forward.switch_mlp.{n}.weight"
+                    ] = mx.stack(to_join)
+        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
+        ]
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if path.endswith("feed_forward.gate"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate
+
+    @property
+    def cast_predicate(self):
+        def predicate(k):
+            return "expert_bias" not in k
+
+        return predicate

mlx_lm/models/llama.py

@@ -1,12 +1,13 @@
 # Copyright © 2023-2024 Apple Inc.
 
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Union
+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 .cache import KVCache, RotatingKVCache
 from .rope_utils import initialize_rope
 
 
@@ -28,11 +29,16 @@ class ModelArgs(BaseModelArgs):
     rope_traditional: bool = False
     rope_scaling: Optional[Dict[str, Union[float, str]]] = None
     tie_word_embeddings: bool = True
+    layer_types: Optional[List[str]] = None
+    sliding_window: Optional[int] = None
 
     def __post_init__(self):
         if self.num_key_value_heads is None:
             self.num_key_value_heads = self.num_attention_heads
 
+        if self.layer_types is None:
+            self.layer_types = ["full_attention"] * self.num_hidden_layers
+
 
 class Attention(nn.Module):
     def __init__(self, args: ModelArgs):
@@ -114,10 +120,11 @@ class MLP(nn.Module):
 
 
 class TransformerBlock(nn.Module):
-    def __init__(self, args: ModelArgs):
+    def __init__(self, args: ModelArgs, use_sliding: bool = False):
         super().__init__()
         self.num_attention_heads = args.num_attention_heads
         self.hidden_size = args.hidden_size
+        self.use_sliding = use_sliding
         self.self_attn = Attention(args)
         self.mlp = MLP(args)
         self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
@@ -145,12 +152,21 @@ class LlamaModel(nn.Module):
         self.args = args
         self.vocab_size = args.vocab_size
         self.num_hidden_layers = args.num_hidden_layers
+        self.layer_types = args.layer_types
+        self.sliding_window = args.sliding_window
         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)
+            TransformerBlock(args=args, use_sliding=layer_type == "sliding_attention")
+            for layer_type in self.layer_types
         ]
         self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.fa_idx = self.layer_types.index("full_attention")
+        self.swa_idx = None
+        for e, l in enumerate(self.layers):
+            if l.use_sliding:
+                self.swa_idx = e
+                break
 
     def __call__(
         self,
@@ -166,10 +182,15 @@ class LlamaModel(nn.Module):
         if cache is None:
             cache = [None] * len(self.layers)
 
-        mask = create_attention_mask(h, cache[0])
+        fa_mask = create_attention_mask(h, cache[self.fa_idx])
+        if self.swa_idx is not None:
+            swa_mask = create_attention_mask(
+                h, cache[self.swa_idx], window_size=self.sliding_window
+            )
 
-        for layer, c in zip(self.layers, cache):
-            h = layer(h, mask, cache=c)
+        for layer, cache in zip(self.layers, cache):
+            mask = swa_mask if layer.use_sliding else fa_mask
+            h = layer(h, mask, cache=cache)
 
         return self.norm(h)
 
@@ -208,3 +229,13 @@ class Model(nn.Module):
     @property
     def layers(self):
         return self.model.layers
+
+    def make_cache(self):
+        return [
+            (
+                RotatingKVCache(max_size=self.model.sliding_window)
+                if layer.use_sliding
+                else KVCache()
+            )
+            for layer in self.layers
+        ]

mlx_lm/models/longcat_flash.py

@@ -235,7 +235,7 @@ class LongcatFlashMoE(nn.Module):
 
         regular_outputs = self.switch_mlp(hidden_states, topk_indices)
 
-        weighted_outputs = regular_outputs * topk_weights[..., None]
+        weighted_outputs = regular_outputs * regular_weights[..., None]
 
         # Add identity expert contribution if needed
         assert self.zero_expert_type == "identity"

mlx_lm/models/minimax.py

@@ -0,0 +1,287 @@
+# 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 .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    intermediate_size: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    max_position_embeddings: int
+    num_experts_per_tok: int
+    num_local_experts: int
+    shared_intermediate_size: int
+    num_hidden_layers: int
+    rms_norm_eps: float
+    rope_theta: float
+    rotary_dim: int
+    vocab_size: int
+    tie_word_embeddings: bool = False
+    scoring_func: str = "sigmoid"
+    head_dim: Optional[int] = None
+    use_qk_norm: bool = True
+
+
+class MiniMaxAttention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        self.hidden_dim = hidden_size = args.hidden_size
+
+        self.num_attention_heads = args.num_attention_heads
+        self.num_key_value_heads = args.num_key_value_heads
+        self.head_dim = head_dim = (
+            args.head_dim or hidden_size // args.num_attention_heads
+        )
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(
+            args.hidden_size, self.num_attention_heads * head_dim, bias=False
+        )
+        self.k_proj = nn.Linear(
+            args.hidden_size, self.num_key_value_heads * head_dim, bias=False
+        )
+        self.v_proj = nn.Linear(
+            args.hidden_size, self.num_key_value_heads * head_dim, bias=False
+        )
+        self.o_proj = nn.Linear(
+            self.num_attention_heads * head_dim, args.hidden_size, bias=False
+        )
+
+        self.use_qk_norm = args.use_qk_norm if hasattr(args, "use_qk_norm") else False
+        if self.use_qk_norm:
+            self.q_norm = nn.RMSNorm(
+                head_dim * self.num_attention_heads, eps=args.rms_norm_eps
+            )
+            self.k_norm = nn.RMSNorm(
+                head_dim * self.num_key_value_heads, eps=args.rms_norm_eps
+            )
+
+        self.rope = nn.RoPE(args.rotary_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)
+
+        if self.use_qk_norm:
+            queries = self.q_norm(queries)
+            keys = self.k_norm(keys)
+
+        queries = queries.reshape(B, L, self.num_attention_heads, -1).transpose(
+            0, 2, 1, 3
+        )
+        keys = 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 MiniMaxSparseMoeBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_experts_per_tok = args.num_experts_per_tok
+
+        self.gate = nn.Linear(args.hidden_size, args.num_local_experts, bias=False)
+        self.switch_mlp = SwitchGLU(
+            args.hidden_size, args.intermediate_size, args.num_local_experts
+        )
+        self.e_score_correction_bias = mx.zeros((args.num_local_experts,))
+
+    def __call__(self, x: mx.array) -> mx.array:
+        gates = self.gate(x.astype(mx.float32))
+
+        scores = mx.sigmoid(gates)
+        orig_scores = scores
+        scores = scores + self.e_score_correction_bias
+
+        k = self.num_experts_per_tok
+        inds = mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k]
+        scores = mx.take_along_axis(orig_scores, inds, axis=-1)
+
+        scores = scores / (mx.sum(scores, axis=-1, keepdims=True) + 1e-20)
+        scores = scores.astype(x.dtype)
+
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2)
+        return y
+
+
+class MiniMaxDecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        self.self_attn = MiniMaxAttention(args)
+
+        self.block_sparse_moe = MiniMaxSparseMoeBlock(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 = x + self.self_attn(self.input_layernorm(x), mask, cache)
+        r = r + self.block_sparse_moe(self.post_attention_layernorm(r))
+        return r
+
+
+class MiniMaxModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+
+        self.layers = [
+            MiniMaxDecoderLayer(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 cache is None:
+            cache = [None] * len(self.layers)
+
+        mask = create_attention_mask(h, cache[0])
+
+        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 = MiniMaxModel(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: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ):
+        out = self.model(inputs=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):
+        """Dequantize FP8 weights and restructure MoE experts."""
+
+        def dequant(weight, scale_inv):
+            dtype = weight.dtype
+            bs = 128  # block size
+            m, n = weight.shape
+            pad_bottom = (-m) % bs
+            pad_side = (-n) % bs
+            weight = mx.pad(weight, ((0, pad_bottom), (0, pad_side)))
+            weight = weight.reshape(
+                ((m + pad_bottom) // bs, bs, (n + pad_side) // bs, bs)
+            )
+            weight = (weight * scale_inv[:, None, :, None]).reshape(
+                m + pad_bottom, n + pad_side
+            )
+            return weight[:m, :n].astype(dtype)
+
+        # Dequantize
+        new_weights = {}
+        for k, v in weights.items():
+            if "weight_scale_inv" in k:
+                scale_inv = v
+                wk = k.replace("_scale_inv", "")
+                weight = weights[wk]
+                weight = dequant(weight, scale_inv)
+                new_weights[wk] = weight
+            elif k not in new_weights:
+                new_weights[k] = v
+        weights = new_weights
+
+        # Step 2: Handle MoE expert weights restructuring
+        if "model.layers.0.block_sparse_moe.experts.0.w1.weight" not in weights:
+            return weights
+
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}"
+            mapping = {"w1": "gate_proj", "w2": "down_proj", "w3": "up_proj"}
+            for orig_name, new_name in mapping.items():
+                if f"{prefix}.block_sparse_moe.experts.0.{orig_name}.weight" in weights:
+                    to_join = [
+                        weights.pop(
+                            f"{prefix}.block_sparse_moe.experts.{e}.{orig_name}.weight"
+                        )
+                        for e in range(self.args.num_local_experts)
+                    ]
+                    weights[
+                        f"{prefix}.block_sparse_moe.switch_mlp.{new_name}.weight"
+                    ] = mx.stack(to_join)
+
+        return weights
+
+    @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
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if path.endswith("block_sparse_moe.gate"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate

mlx_lm/models/ministral3.py

@@ -0,0 +1,264 @@
+# Copyright © 2023-2024 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 .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
+    head_dim: Optional[int] = None
+    max_position_embeddings: Optional[int] = None
+    num_key_value_heads: Optional[int] = None
+    rope_parameters: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = True
+    layer_types: Optional[List[str]] = None
+    sliding_window: Optional[int] = None
+
+    def __post_init__(self):
+        if self.num_key_value_heads is None:
+            self.num_key_value_heads = self.num_attention_heads
+
+        if self.layer_types is None:
+            self.layer_types = ["full_attention"] * self.num_hidden_layers
+
+
+def _get_llama_4_attn_scale(
+    start: int, stop: int, beta: float, max_position_embeddings: int
+):
+    scaling = 1 + beta * mx.log(
+        1 + mx.floor(mx.arange(start, stop) / max_position_embeddings)
+    )
+    return scaling[:, 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
+        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
+
+        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.rope = initialize_rope(
+            self.head_dim,
+            args.rope_parameters["rope_theta"],
+            False,
+            args.rope_parameters,
+            args.max_position_embeddings,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        attn_scale: 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)
+
+        offset = 0
+        if cache is not None:
+            offset = cache.offset
+            queries = self.rope(queries, offset=offset)
+            keys = self.rope(keys, offset=offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+        queries = queries * attn_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.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, use_sliding: bool = False):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.hidden_size = args.hidden_size
+        self.use_sliding = use_sliding
+        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,
+        attn_scale: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), attn_scale, mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        out = h + r
+        return out
+
+
+class LanguageModel(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.layer_types = args.layer_types
+        self.sliding_window = args.sliding_window
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            TransformerBlock(args=args, use_sliding=layer_type == "sliding_attention")
+            for layer_type in self.layer_types
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.fa_idx = self.layer_types.index("full_attention")
+        self.swa_idx = None
+        for e, l in enumerate(self.layers):
+            if l.use_sliding:
+                self.swa_idx = e
+                break
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        if input_embeddings is not None:
+            h = input_embeddings
+        else:
+            h = self.embed_tokens(inputs)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+            offset = 0
+        else:
+            offset = cache[0].offset
+
+        fa_mask = create_attention_mask(h, cache[self.fa_idx])
+        if self.swa_idx is not None:
+            swa_mask = create_attention_mask(
+                h, cache[self.swa_idx], window_size=self.sliding_window
+            )
+
+        attn_scale = _get_llama_4_attn_scale(
+            offset,
+            offset + inputs.shape[1],
+            self.args.rope_parameters["llama_4_scaling_beta"],
+            self.args.rope_parameters["original_max_position_embeddings"],
+        ).astype(h.dtype)
+
+        for layer, cache in zip(self.layers, cache):
+            mask = swa_mask if layer.use_sliding else fa_mask
+            h = layer(h, attn_scale, mask, cache=cache)
+
+        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 = LanguageModel(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        out = self.model(inputs, 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
+
+    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)
+
+        new_weights = {}
+        for k, v in weights.items():
+            if "weight_scale_inv" in k:
+                scale_inv = v
+                wk = k.replace("_scale_inv", "")
+                weight = weights[wk]
+                new_weights[wk] = weight * scale_inv
+            elif "activation_scale" in k:
+                continue
+            elif k not in new_weights:
+                new_weights[k] = v
+        weights = new_weights
+
+        return weights
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def make_cache(self):
+        return [
+            (
+                RotatingKVCache(max_size=self.model.sliding_window)
+                if layer.use_sliding
+                else KVCache()
+            )
+            for layer in self.layers
+        ]

mlx_lm/models/mistral3.py

@@ -7,7 +7,7 @@ import mlx.core as mx
 import mlx.nn as nn
 from mlx.utils import tree_flatten, tree_unflatten
 
-from . import llama
+from . import llama, ministral3
 from .base import BaseModelArgs
 
 
@@ -17,7 +17,8 @@ class ModelArgs(BaseModelArgs):
     text_config: dict
 
     def __post_init__(self):
-        self.text_config["tie_word_embeddings"] = False
+        if "tie_word_embeddings" not in self.text_config:
+            self.text_config["tie_word_embeddings"] = False
 
 
 class Model(nn.Module):
@@ -25,7 +26,14 @@ class Model(nn.Module):
         super().__init__()
         self.args = args
         self.model_type = args.model_type
-        self.language_model = llama.Model(llama.ModelArgs.from_dict(args.text_config))
+        if args.text_config.get("model_type") == "ministral3":
+            self.language_model = ministral3.Model(
+                ministral3.ModelArgs.from_dict(args.text_config)
+            )
+        else:
+            self.language_model = llama.Model(
+                llama.ModelArgs.from_dict(args.text_config)
+            )
 
     def __call__(
         self,
@@ -41,6 +49,8 @@ class Model(nn.Module):
         weights = tree_unflatten(list(weights.items()))
         weights.pop("vision_tower", None)
         weights.pop("multi_modal_projector", None)
+        lm_weights = dict(tree_flatten(weights["language_model"]))
+        weights["language_model"] = self.language_model.sanitize(lm_weights)
         return dict(tree_flatten(weights))
 
     @property

mlx_lm/models/nanochat.py

@@ -0,0 +1,233 @@
+# 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_attention_mask, scaled_dot_product_attention
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str = "nanochat"
+    hidden_size: int = 1280
+    num_hidden_layers: int = 20
+    num_attention_heads: int = 10
+    num_key_value_heads: int = 10
+    vocab_size: int = 65536
+    max_position_embeddings: int = 2048
+    intermediate_size: int = 5120  # 4 * hidden_size
+    rope_theta: float = 10000.0
+
+
+def rms_norm(x):
+    """Functional RMSNorm with no learnable parameters."""
+    return mx.fast.rms_norm(x, None, 1e-5)
+
+
+def apply_rotary_emb(x, offset, base=10000.0, freqs=None):
+    """Apply RoPE with blocked layout.
+
+
+    Args:
+        x: Input tensor in (B, H, T, D) format
+        offset: Position offset for KV caching
+        base: RoPE base frequency (default 10000.0)
+        freqs: Precomputed negated frequencies (optional)
+
+    Returns:
+        Tensor with RoPE applied, same shape as input
+    """
+    head_dim = x.shape[-1]
+
+    if freqs is None:
+        # Compute negated frequencies
+        half_D = head_dim // 2
+        freqs = -mx.exp(
+            mx.arange(0.0, half_D, dtype=mx.float32) * (math.log(base) / half_D)
+        )
+
+    # Use traditional=False + negated freqs
+    return mx.fast.rope(
+        x,
+        dims=head_dim,
+        traditional=False,
+        base=None,
+        freqs=freqs,
+        scale=1.0,
+        offset=offset,
+    )
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        self.hidden_size = args.hidden_size
+        self.num_heads = args.num_attention_heads
+        self.num_kv_heads = args.num_key_value_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.scale = self.head_dim**-0.5
+        self.rope_theta = args.rope_theta
+
+        self.c_q = nn.Linear(
+            self.hidden_size, self.num_heads * self.head_dim, bias=False
+        )
+        self.c_k = nn.Linear(
+            self.hidden_size, self.num_kv_heads * self.head_dim, bias=False
+        )
+        self.c_v = nn.Linear(
+            self.hidden_size, self.num_kv_heads * self.head_dim, bias=False
+        )
+        self.c_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+
+        # Precompute negated RoPE frequencies for awni's approach
+        half_D = self.head_dim // 2
+        self._rope_freqs = -mx.exp(
+            mx.arange(0.0, half_D, dtype=mx.float32)
+            * (math.log(self.rope_theta) / half_D)
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, _ = x.shape
+
+        queries = self.c_q(x)
+        keys = self.c_k(x)
+        values = self.c_v(x)
+
+        # Reshape to (B, L, H, D) then transpose to (B, H, L, D)
+        queries = queries.reshape(B, L, self.num_heads, self.head_dim).transpose(
+            0, 2, 1, 3
+        )
+        keys = keys.reshape(B, L, self.num_kv_heads, self.head_dim).transpose(
+            0, 2, 1, 3
+        )
+        values = values.reshape(B, L, self.num_kv_heads, self.head_dim).transpose(
+            0, 2, 1, 3
+        )
+
+        # Apply RoPE using precomputed frequencies (expects B, H, T, D format)
+        offset = cache.offset if cache is not None else 0
+        queries = apply_rotary_emb(
+            queries, offset=offset, base=self.rope_theta, freqs=self._rope_freqs
+        )
+        keys = apply_rotary_emb(
+            keys, offset=offset, base=self.rope_theta, freqs=self._rope_freqs
+        )
+
+        # QK norm (critical feature of nanochat!)
+        queries = rms_norm(queries)
+        keys = rms_norm(keys)
+
+        # Handle KV cache after transpose
+        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
+        )
+
+        # Reshape back
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, self.hidden_size)
+        return self.c_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.c_fc = nn.Linear(args.hidden_size, args.intermediate_size, bias=False)
+        self.c_proj = nn.Linear(args.intermediate_size, args.hidden_size, bias=False)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        # Critical: nanochat uses ReLU^2, not GELU!
+        x = self.c_fc(x)
+        x = nn.relu2(x)
+        return self.c_proj(x)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.attn = Attention(args)
+        self.mlp = MLP(args)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        # Pre-norm architecture with functional RMSNorm
+        h = x + self.attn(rms_norm(x), mask=mask, cache=cache)
+        out = h + self.mlp(rms_norm(h))
+        return out
+
+
+class NanoChatModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.wte = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.h = [TransformerBlock(args) for _ in range(args.num_hidden_layers)]
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+    ) -> mx.array:
+        h = self.wte(inputs)
+        # Critical: norm after token embedding
+        h = rms_norm(h)
+
+        if cache is None:
+            cache = [None] * len(self.h)
+
+        mask = create_attention_mask(h, cache[0])
+
+        for layer, c in zip(self.h, cache):
+            h = layer(h, mask=mask, cache=c)
+
+        # Critical: final norm before lm_head
+        h = rms_norm(h)
+
+        return h
+
+
+@partial(mx.compile, shapeless=True)
+def softcap(logits, cap=15.0):
+    return cap * mx.tanh(logits / cap)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.transformer = NanoChatModel(args)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+    ) -> mx.array:
+        out = self.transformer(inputs, cache=cache)
+        logits = self.lm_head(out)
+
+        # Critical: logits softcap (nanochat uses softcap=15)
+        logits = softcap(logits)
+
+        return logits
+
+    @property
+    def layers(self):
+        return self.transformer.h

mlx_lm/models/olmo3.py

@@ -0,0 +1,235 @@
+# 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 .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
+    max_position_embeddings: int
+    sliding_window: int
+    rope_theta: float
+    attention_bias: bool = False
+    layer_types: Optional[List[str]] = None
+    num_key_value_heads: Optional[int] = None
+    head_dim: Optional[int] = None
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = False
+
+    def __post_init__(self):
+        if self.num_key_value_heads is None:
+            self.num_key_value_heads = self.num_attention_heads
+        if self.layer_types is None:
+            self.layer_types = [
+                "full_attention" if (i + 1) % 4 == 0 else "sliding_attention"
+                for i in range(self.num_hidden_layers)
+            ]
+
+
+class Olmo3Attention(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.num_key_value_heads = args.num_key_value_heads
+        self.layer_idx = layer_idx
+
+        self.head_dim = args.head_dim or 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=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=args.attention_bias,
+        )
+
+        self.q_norm = nn.RMSNorm(
+            args.num_attention_heads * self.head_dim, eps=args.rms_norm_eps
+        )
+        self.k_norm = nn.RMSNorm(
+            args.num_key_value_heads * self.head_dim, eps=args.rms_norm_eps
+        )
+
+        if args.layer_types[layer_idx] != "full_attention":
+            self.rope = nn.RoPE(self.head_dim, traditional=False, base=args.rope_theta)
+        else:
+            self.rope = initialize_rope(
+                self.head_dim,
+                traditional=False,
+                base=args.rope_theta,
+                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, _ = x.shape
+        queries = self.q_norm(self.q_proj(x))
+        keys = self.k_norm(self.k_proj(x))
+        values = self.v_proj(x)
+
+        queries = queries.reshape(B, L, self.num_attention_heads, -1).transpose(
+            0, 2, 1, 3
+        )
+        keys = 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 Olmo3MLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.gate_proj = nn.Linear(args.hidden_size, args.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(args.intermediate_size, args.hidden_size, bias=False)
+        self.up_proj = nn.Linear(args.hidden_size, args.intermediate_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 Olmo3DecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.hidden_size = args.hidden_size
+        self.self_attn = Olmo3Attention(args, layer_idx=layer_idx)
+        self.mlp = Olmo3MLP(args)
+        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.post_attention_layernorm(self.self_attn(x, mask, cache))
+        h = x + r
+        r = self.post_feedforward_layernorm(self.mlp(h))
+        out = h + r
+        return out
+
+
+class Olmo3Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.sliding_window = args.sliding_window
+
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            Olmo3DecoderLayer(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)
+
+        self.swa_idx = args.layer_types.index("sliding_attention")
+        self.ga_idx = args.layer_types.index("full_attention")
+        self.layer_types = args.layer_types
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        h = self.embed_tokens(inputs)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        full_mask = create_attention_mask(h, cache[self.ga_idx])
+        sliding_window_mask = create_attention_mask(
+            h, cache[self.swa_idx], window_size=self.sliding_window
+        )
+
+        for layer, c, layer_type in zip(self.layers, cache, self.layer_types):
+            mask = full_mask if layer_type == "full_attention" else sliding_window_mask
+            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 = Olmo3Model(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        out = self.model(inputs, cache)
+        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 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))
+        return caches

mlx_lm/models/pipeline.py

@@ -0,0 +1,31 @@
+# Copyright © 2025 Apple Inc.
+
+import mlx.core as mx
+
+
+class PipelineMixin:
+    def __init__(self):
+        super().__init__()
+        self.pipeline_rank = 0
+        self.pipeline_size = 1
+        self.start_idx = 0
+        self.end_idx = None
+
+    @property
+    def pipeline_layers(self):
+        return self.layers[self.start_idx : self.end_idx]
+
+    def pipeline(self, group):
+        # Split layers in reverse so rank=0 gets the last layers and
+        # rank=pipeline_size-1 gets the first
+        self.pipeline_rank = group.rank()
+        self.pipeline_size = group.size()
+        layers_per_rank = len(self.layers) // self.pipeline_size
+        extra = len(self.layers) - layers_per_rank * self.pipeline_size
+        if self.pipeline_rank < extra:
+            layers_per_rank += 1
+        self.start_idx = (self.pipeline_size - self.pipeline_rank - 1) * layers_per_rank
+        self.end_idx = self.start_idx + layers_per_rank
+        self.layers = self.layers[: self.end_idx]
+        # Keep the layer numbers the same for model loading
+        self.layers[: self.start_idx] = [None] * self.start_idx

mlx_lm/models/qwen2_vl.py

@@ -20,7 +20,7 @@ class ModelArgs(BaseModelArgs):
     def from_dict(cls, params):
         if "text_config" not in params:
             return cls(model_type=params["model_type"], text_config=params)
-        return cls(**params)
+        return super().from_dict(params)
 
 
 class Model(nn.Module):

mlx_lm/models/qwen3.py

@@ -141,8 +141,12 @@ class Qwen3Model(nn.Module):
         self,
         inputs: mx.array,
         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 cache is None:
             cache = [None] * len(self.layers)
@@ -167,8 +171,9 @@ class Model(nn.Module):
         self,
         inputs: mx.array,
         cache=None,
+        input_embeddings: Optional[mx.array] = None,
     ):
-        out = self.model(inputs, cache)
+        out = self.model(inputs, cache, input_embeddings)
         if self.args.tie_word_embeddings:
             out = self.model.embed_tokens.as_linear(out)
         else:

mlx_lm/models/qwen3_moe.py

@@ -188,8 +188,12 @@ class Qwen3MoeModel(nn.Module):
         self,
         inputs: mx.array,
         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 cache is None:
             cache = [None] * len(self.layers)
@@ -211,11 +215,9 @@ class Model(nn.Module):
         self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
 
     def __call__(
-        self,
-        inputs: mx.array,
-        cache=None,
+        self, inputs: mx.array, cache=None, input_embeddings: Optional[mx.array] = None
     ):
-        out = self.model(inputs, cache)
+        out = self.model(inputs, cache, input_embeddings)
         return self.lm_head(out)
 
     def sanitize(self, weights):

mlx_lm/models/qwen3_next.py

@@ -1,12 +1,19 @@
 # Copyright © 2025 Apple Inc.
 
+from __future__ import annotations
+
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .base import (
+    BaseModelArgs,
+    create_attention_mask,
+    create_ssm_mask,
+    scaled_dot_product_attention,
+)
 from .cache import KVCache, MambaCache
 from .gated_delta import gated_delta_update
 from .rope_utils import initialize_rope
@@ -237,6 +244,8 @@ class Qwen3NextGatedDeltaNet(nn.Module):
         mixed_qkv = mx.concatenate(
             [q.reshape(B, S, -1), k.reshape(B, S, -1), v.reshape(B, S, -1)], axis=-1
         )
+        if mask is not None:
+            mixed_qkv = mx.where(mask[..., None], mixed_qkv, 0)
         conv_input = mx.concatenate([conv_state, mixed_qkv], axis=1)
         if cache is not None:
             cache[0] = conv_input[:, -(self.conv_kernel_size - 1) :]
@@ -251,14 +260,23 @@ class Qwen3NextGatedDeltaNet(nn.Module):
             )
         ]
 
-        if cache is not None:
-            state = cache[1]
-
+        state = cache[1] if cache else None
         inv_scale = k.shape[-1] ** -0.5
         q = (inv_scale**2) * mx.fast.rms_norm(q, None, 1e-6)
         k = inv_scale * mx.fast.rms_norm(k, None, 1e-6)
 
-        out, state = gated_delta_update(q, k, v, a, b, self.A_log, self.dt_bias, state)
+        out, state = gated_delta_update(
+            q,
+            k,
+            v,
+            a,
+            b,
+            self.A_log,
+            self.dt_bias,
+            state,
+            mask,
+            use_kernel=not self.training,
+        )
 
         if cache is not None:
             cache[1] = state
@@ -350,6 +368,7 @@ class Qwen3NextModel(nn.Module):
             for i in range(args.num_hidden_layers)
         ]
         self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.ssm_idx = 0
         self.fa_idx = args.full_attention_interval - 1
 
     def __call__(
@@ -362,9 +381,11 @@ class Qwen3NextModel(nn.Module):
         if cache is None:
             cache = [None] * len(self.layers)
 
-        mask = create_attention_mask(hidden_states, cache[self.fa_idx])
+        fa_mask = create_attention_mask(hidden_states, cache[self.fa_idx])
+        ssm_mask = create_ssm_mask(hidden_states, cache[self.ssm_idx])
 
         for layer, c in zip(self.layers, cache):
+            mask = ssm_mask if layer.is_linear else fa_mask
             hidden_states = layer(hidden_states, mask=mask, cache=c)
 
         return self.norm(hidden_states)

mlx_lm/models/qwen3_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 qwen3
+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 super().from_dict(params)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.language_model = qwen3.Model(qwen3.ModelArgs.from_dict(args.text_config))
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        return self.language_model(
+            inputs, cache=cache, input_embeddings=input_embeddings
+        )
+
+    def sanitize(self, weights):
+        weights = tree_unflatten(list(weights.items()))
+        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_vl_moe.py

@@ -0,0 +1,77 @@
+# 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 qwen3_moe
+from .base import BaseModelArgs
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    text_config: dict
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.language_model = qwen3_moe.Model(
+            qwen3_moe.ModelArgs.from_dict(args.text_config)
+        )
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        return self.language_model(
+            inputs, cache=cache, input_embeddings=input_embeddings
+        )
+
+    def sanitize(self, weights):
+        weights = tree_unflatten(list(weights.items()))
+        weights.pop("visual", None)
+        weights = dict(
+            tree_flatten(
+                {
+                    "language_model": {
+                        "model": weights["language_model"]["model"],
+                        "lm_head": weights["language_model"]["lm_head"],
+                    }
+                }
+            )
+        )
+
+        for l in range(self.language_model.args.num_hidden_layers):
+            prefix = f"language_model.model.layers.{l}.mlp"
+            gate_up_key = f"{prefix}.experts.gate_up_proj"
+            if gate_up_key in weights:
+                gate_up = weights.pop(gate_up_key)
+                mid = gate_up.shape[-1] // 2
+                weights[f"{prefix}.switch_mlp.gate_proj.weight"] = gate_up[
+                    ..., :mid
+                ].swapaxes(-2, -1)
+                weights[f"{prefix}.switch_mlp.up_proj.weight"] = gate_up[
+                    ..., mid:
+                ].swapaxes(-2, -1)
+                weights[f"{prefix}.switch_mlp.down_proj.weight"] = weights.pop(
+                    f"{prefix}.experts.down_proj"
+                ).swapaxes(-2, -1)
+
+        return weights
+
+    @property
+    def quant_predicate(self):
+        return self.language_model.quant_predicate
+
+    @property
+    def layers(self):
+        return self.language_model.model.layers

mlx_lm/models/rope_utils.py

@@ -168,9 +168,7 @@ class YarnRoPE(nn.Module):
             scaling_factor, mscale_all_dim
         )
         freq_extra = base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
-        freq_inter = scaling_factor * base ** (
-            mx.arange(0, dims, 2, dtype=mx.float32) / dims
-        )
+        freq_inter = scaling_factor * freq_extra
         low, high = yarn_find_correction_range()
         freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dims // 2)
         self._freqs = (freq_inter * freq_extra) / (

mlx_lm/models/ssm.py

@@ -114,6 +114,7 @@ def ssm_attn(
     state: Optional[mx.array] = None,
     time_step_limit: Tuple[float, float] = (0.001, 100.0),
     mask: Optional[mx.array] = None,
+    step: int = 256,
 ) -> Tuple[mx.array, mx.array]:
     """SSD-SSM forward pass.
 
@@ -127,6 +128,7 @@ def ssm_attn(
         dt_bias: Bias for time deltas of shape (num_heads,).
         time_step_limit: Minimum and maximum value for time deltas.
         mask: Optional multiplicative mask.
+        step: Step size for processing x.
 
     Code modified from
     https://github.com/cartesia-ai/edge/blob/main/cartesia-mlx/cartesia_mlx/layers/ssd/ops.py
@@ -138,38 +140,52 @@ def ssm_attn(
     dt = compute_dt(dt, dt_bias, time_step_limit)
     repeats = h // g
     A = -mx.exp(A_log)
-    B = mx.transpose(B, (0, 2, 3, 1))
-
-    # A * s + B * C
-    CB = mx.swapaxes(C, 1, 2) @ B
-    CB = mx.repeat(CB, repeats, axis=1)
-
     dtA = dt * A.reshape(1, 1, -1)
-    decay = mx.exp(segsum(dtA.swapaxes(1, 2), mask=mask))
-
-    surrogate_attention_matrix = mx.tril(CB * decay, 0)
-
     dtx = dt.reshape(b, l, h, 1) * x
-    y = surrogate_attention_matrix @ dtx.swapaxes(1, 2)
-    y = mx.swapaxes(y, 1, 2)
 
-    decay = decay[:, :, -1:, :].transpose(0, 3, 1, 2)
-    B = mx.repeat(B, h // g, axis=1).swapaxes(2, 3)
-    dtxdecay = dtx * decay
-    dtxdecay = dtxdecay.swapaxes(1, 2).swapaxes(2, 3)
+    def _step(dtx, dtA, B, C, state, mask):
+        s = dtx.shape[1]
+        B = mx.transpose(B, (0, 2, 3, 1))
 
-    next_state = dtxdecay @ B
+        CB = mx.swapaxes(C, 1, 2) @ B
+        CB = mx.repeat(CB, repeats, axis=1)
 
-    if state is not None:
-        exp_dtA_cumsum = mx.exp(mx.cumsum(dtA, axis=-2))
-        next_state += exp_dtA_cumsum[:, -1, :, None, None] * state
-        state = state.reshape((b, 1, g, repeats, dh, d))
-        C = C.reshape(b, l, g, 1, d, 1)
-        y_prev = (state @ C).squeeze(-1).flatten(2, 3)
-        y += exp_dtA_cumsum[..., None] * y_prev
+        decay = mx.exp(segsum(dtA.swapaxes(1, 2), mask=mask))
 
-    y += x * D.reshape(1, 1, h, 1)
-    return y, next_state
+        surrogate_attention_matrix = mx.tril(CB * decay, 0)
+
+        y = surrogate_attention_matrix @ dtx.swapaxes(1, 2)
+        y = mx.swapaxes(y, 1, 2)
+
+        decay = decay[:, :, -1:, :].transpose(0, 3, 1, 2)
+        B = mx.repeat(B, h // g, axis=1).swapaxes(2, 3)
+        dtxdecay = dtx * decay
+        dtxdecay = dtxdecay.swapaxes(1, 2).swapaxes(2, 3)
+
+        next_state = dtxdecay @ B
+
+        if state is not None:
+            exp_dtA_cumsum = mx.exp(mx.cumsum(dtA, axis=-2))
+            next_state += exp_dtA_cumsum[:, -1, :, None, None] * state
+            state = state.reshape((b, 1, g, repeats, dh, d))
+            C = C.reshape(b, s, g, 1, d, 1)
+            y_prev = (state @ C).squeeze(-1).flatten(2, 3)
+            y += exp_dtA_cumsum[..., None] * y_prev
+        return y, next_state
+
+    ys = []
+    for i in range(0, l, step):
+        y, state = _step(
+            dtx[:, i : i + step],
+            dtA[:, i : i + step],
+            B[:, i : i + step],
+            C[:, i : i + step],
+            state,
+            None if mask is None else mask[..., i : i + step],
+        )
+        ys.append(y)
+    y = mx.concatenate(ys, axis=1) + x * D.reshape(1, 1, h, 1)
+    return y, state
 
 
 def ssm_update(

mlx_lm/perplexity.py

@@ -13,8 +13,7 @@ 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
+from mlx_lm.utils import get_total_parameters, load
 
 
 def load_data(

mlx_lm/quant/awq.py

@@ -16,8 +16,7 @@ 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,
+    load,
     save,
 )
 
@@ -555,8 +554,7 @@ def main():
 
     mx.random.seed(args.seed)
 
-    model_path, hf_repo = get_model_path(args.model, revision=None)
-    model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
+    model, tokenizer, config = load(args.model, lazy=True, return_config=True)
 
     model_type = config["model_type"]
     if (awq_config := AWQ_MODEL_CONFIGS.get(model_type, None)) is None:
@@ -580,9 +578,8 @@ def main():
     config = update_config(model, config)
     save(
         args.mlx_path,
-        model_path,
+        args.model,
         model,
         tokenizer,
         config,
-        hf_repo=hf_repo,
     )

mlx_lm/quant/dwq.py

@@ -4,12 +4,13 @@ import argparse
 import copy
 import time
 import types
+from pathlib import Path
 
 import mlx.core as mx
 import mlx.nn as nn
 import mlx.optimizers as optimizers
 import numpy as np
-from mlx.utils import tree_flatten, tree_map
+from mlx.utils import tree_map
 from tqdm import tqdm
 
 from mlx_lm.tuner.datasets import load_dataset
@@ -17,44 +18,66 @@ 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,
+    load,
+    load_tokenizer,
+    pipeline_load,
     quantize_model,
     save,
 )
 
 
-class Catcher(nn.Module):
-    def __init__(self, module):
-        super().__init__()
-        self.module = module
+def compute_dwq_targets(
+    model,
+    save_dir,
+    train_data,
+    valid_data,
+    batch_size,
+    max_seq_length,
+    seed,
+):
+    rank = mx.distributed.init().rank()
 
-    def __call__(self, *args, **kwargs):
-        outputs = self.module(*args, **kwargs)
-        self.outputs = outputs[0] if isinstance(outputs, tuple) else outputs
-        return outputs
+    def _compute_targets(data, path, split):
 
-    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)
+        if rank == 0:
+            path = path / split
+            path.mkdir(parents=True, exist_ok=True)
+        for i, (batch, _) in (
+            pbar := tqdm(
+                enumerate(iterate_batches(data, batch_size, max_seq_length, seed=seed)),
+                total=len(data) // batch_size,
+                desc=f"Computing targets for {split}",
+                disable=rank != 0,
+            )
+        ):
+            batch = batch[:, :-1]
+            logits = model(batch)
+            # Hack to make the last op pre-eval on the CPU to avoid even timeout
+            logits = mx.stop_gradient(logits, stream=mx.cpu)
+            mx.eval(logits)
+            if rank == 0:
+                idx = mx.argpartition(logits, kth=-1024, axis=-1)[..., -1024:]
+                logits = mx.take_along_axis(logits, idx, axis=-1)
+
+                file = path / f"{i:010d}.safetensors"
+                mx.save_safetensors(file, {"logits": logits, "indices": idx})
+
+    _compute_targets(valid_data, save_dir, "valid")
+    _compute_targets(train_data, save_dir, "train")
 
 
 def dwq_quantize(
     model,
-    q_model,
+    target_fn,
     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,
+    batch_size,
+    max_seq_length,
+    seed,
     dtype: mx.Dtype = mx.bfloat16,
     gradient_checkpoint: bool = False,
+    temperature: float = 2.0,
 ):
     group = mx.distributed.init()
     world_size = group.size()
@@ -65,52 +88,38 @@ def dwq_quantize(
             tqdm.write(*args, **kwargs)
 
     def unfreeze(_, m):
-        if hasattr(m, "bits") and hasattr(m, "group_size") and m.mode == "affine":
+        if (
+            hasattr(m, "bits")
+            and hasattr(m, "group_size")
+            and m.mode == "affine"
+            and m.bits < 8
+        ):
             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])
+    model.train()
+    model.apply_to_modules(unfreeze)
+    print_trainable_parameters(model)
 
     if gradient_checkpoint:
-        grad_checkpoint(q_model.layers[0])
+        grad_checkpoint(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
+    scale = 1 / temperature
 
-    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)
+    def loss_fn(params, x, targets, lengths):
+        model.update(tree_map(lambda x: x.astype(dtype), params))
+        logits = model(x)
+        if isinstance(targets, tuple):
+            targets, ids = targets
+            logits = mx.take_along_axis(logits, ids, axis=-1)
+        losses = kl_div_loss(scale * logits, scale * 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
+        loss = (mask * losses).sum() / ntoks
+        return loss, ntoks
 
-    def step(inputs, targets, extra_targets, lengths, params):
-        (loss, ntoks, *_), grads = mx.value_and_grad(loss_fn)(
-            params, inputs, targets, extra_targets, lengths
+    def step(inputs, targets, lengths, params):
+        (loss, ntoks), grads = mx.value_and_grad(loss_fn)(
+            params, inputs, targets, lengths
         )
         grads = nn.average_gradients(grads)
         params = opt.apply_gradients(grads, params)
@@ -118,42 +127,32 @@ def dwq_quantize(
 
     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)),
+        for i, (batch, lengths) in tqdm(
+            enumerate(
+                iterate_batches(valid_data, batch_size, max_seq_length, seed=seed)
+            ),
             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
-            )
+            targets = target_fn(batch, i, split="valid")
+            mx.eval(targets)
+            loss, ntoks = loss_fn(params, batch, 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}")
+        rprint(f"Validation: {it=}, {loss=:.3f}")
         return loss
 
     # Accumulate learned weights in higher precision
     params = tree_map(
         lambda x: x.astype(mx.float32),
-        q_model.trainable_parameters(),
+        model.trainable_parameters(),
     )
 
     total_loss = 0.0
@@ -167,14 +166,16 @@ def dwq_quantize(
 
     for it, (batch, lengths) in (
         pbar := tqdm(
-            enumerate(iterate_batches(train_data, batch_size, max_seq_length)),
+            enumerate(
+                iterate_batches(train_data, batch_size, max_seq_length, seed=seed)
+            ),
             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)
+        targets = target_fn(batch, it, split="train")
+        mx.eval(targets)
+        loss, ntoks, params = step(batch, targets, lengths, params)
         mx.eval(loss, params)
         loss = mx.distributed.all_sum(loss, stream=mx.cpu).item() / world_size
         ntoks = mx.distributed.all_sum(ntoks, stream=mx.cpu).item()
@@ -205,9 +206,7 @@ def dwq_quantize(
             " 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
+    model.update(tree_map(lambda x: x.astype(dtype), params))
 
 
 def load_data(
@@ -248,10 +247,12 @@ def main():
         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`.",
+        type=str,
         required=True,
     )
     parser.add_argument(
         "--quantized-model",
+        type=str,
         default=None,
         help="An already quantized model (the student model) to improve with DWQ.",
     )
@@ -273,7 +274,7 @@ def main():
         default=2048,
         help="Number of samples to use for training.",
     )
-    parser.add_argument("--max-seq-length", type=int, default=2049)
+    parser.add_argument("--max-seq-length", type=int, default=1025)
     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)
@@ -288,30 +289,83 @@ def main():
         action="store_true",
         help="Use gradient checkpointing to reduce memory use.",
     )
+    parser.add_argument(
+        "--target-dir", type=str, default=None, help="Directory to save/load targets."
+    )
+    parser.add_argument(
+        "--targets-only", action="store_true", help="Compute the targets and exit."
+    )
+    parser.add_argument(
+        "--pipeline",
+        action="store_true",
+        help="Use pipeline parallel instead of data parallel.",
+    )
+
     args = parser.parse_args()
 
     group = mx.distributed.init()
 
     num_samples = args.num_samples
-    if num_samples % group.size() > 0:
+    if not args.pipeline and 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
-    )
+    if args.target_dir is not None:
+        target_dir = Path(args.target_dir)
+        has_targets = target_dir.exists()
+    else:
+        has_targets = False
+        target_dir = None
+
+    tokenizer = load_tokenizer(args.model)
 
     train_data, valid_data = load_data(
         tokenizer, args.data_path, args.num_samples, args.max_seq_length
     )
 
+    # Load the base model if we need it
+    if not has_targets or args.quantized_model is None:
+        if args.pipeline and group.size() > 1:
+            model, _, config = pipeline_load(args.model, return_config=True)
+        else:
+            model, _, config = load(args.model, return_config=True, lazy=True)
+    else:
+        model = None
+
+    # Pre-compute the targets
+    if not has_targets and target_dir is not None:
+        compute_dwq_targets(
+            model,
+            target_dir,
+            train_data,
+            valid_data,
+            batch_size=args.batch_size,
+            max_seq_length=args.max_seq_length,
+            seed=args.seed,
+        )
+        has_targets = True
+
+    if args.targets_only:
+        exit(0)
+
+    if has_targets:
+
+        def target_fn(_, idx, split):
+            targets = mx.load(target_dir / split / f"{idx:010d}.safetensors")
+            return targets["logits"], targets["indices"]
+
+    else:
+
+        def target_fn(batch, idx, split):
+            return model(batch)
+
     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
+        q_model, tokenizer, config = load(
+            args.quantized_model,
+            lazy=True,
+            return_config=True,
         )
         if "quantization" not in config:
             raise ValueError("Quantized model must already be quantized.")
@@ -324,22 +378,30 @@ def main():
             bits=args.bits,
         )
 
+    # Delete the base model if it's not needed
+    if has_targets and model is not None:
+        del model
+
+    if mx.metal.is_available():
+        max_rec_size = mx.metal.device_info()["max_recommended_working_set_size"]
+        mx.set_wired_limit(max_rec_size)
+
     opt = optimizers.Adam(learning_rate=args.learning_rate, bias_correction=True)
     dwq_quantize(
-        model,
         q_model,
+        target_fn,
         opt,
         train_data,
         valid_data,
         batch_size=args.batch_size,
         max_seq_length=args.max_seq_length,
+        seed=args.seed,
         gradient_checkpoint=args.grad_checkpoint,
     )
     save(
         args.mlx_path,
-        model_path,
+        args.model,
         q_model,
         tokenizer,
         config,
-        hf_repo=hf_repo,
     )

mlx_lm/quant/dynamic_quant.py

@@ -16,8 +16,6 @@ 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,
@@ -187,9 +185,9 @@ def main():
     args = parser.parse_args()
 
     group = mx.distributed.init()
+    model, tokenizer, config = load(args.model, return_config=True)
 
     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)
 
@@ -211,8 +209,6 @@ def main():
             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)
 
@@ -230,7 +226,7 @@ def main():
         high_group_size=args.high_group_size,
     )
 
-    def quant_predicate(p, m, _):
+    def quant_predicate(p, m):
         if not hasattr(m, "to_quantized"):
             return False
         if sensitivities[p] > threshold:
@@ -251,11 +247,10 @@ def main():
 
     save(
         args.mlx_path,
-        model_path,
+        args.model,
         model,
         tokenizer,
         config,
-        hf_repo=hf_repo,
     )
     print(f"Peak memory used: {mx.get_peak_memory() / 1000**3:.3f}GB")
 

mlx_lm/quant/gptq.py

@@ -18,8 +18,7 @@ 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,
+    load,
     save,
 )
 
@@ -202,8 +201,7 @@ def main():
 
     mx.random.seed(args.seed)
 
-    model_path, hf_repo = get_model_path(args.model, revision=None)
-    model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
+    model, tokenizer, config = load(args.model, lazy=True, return_config=True)
     calibration_data = load_data(tokenizer, args.num_samples, args.sequence_length)
 
     model, config["quantization"] = gptq_quantize(
@@ -220,11 +218,10 @@ def main():
 
     save(
         args.mlx_path,
-        model_path,
+        args.model,
         model,
         tokenizer,
         config,
-        hf_repo=hf_repo,
     )
 
 

mlx_lm/tokenizer_utils.py

@@ -1,7 +1,7 @@
 import json
 from functools import partial
 from json import JSONDecodeError
-from typing import List
+from typing import Any, Dict, List, Optional
 
 from transformers import AutoTokenizer, PreTrainedTokenizerFast
 
@@ -210,7 +210,7 @@ class BPEStreamingDetokenizer(StreamingDetokenizer):
         # For multi-byte utf-8 wait until they are complete
         # For single spaces wait until the next token to clean it if needed
         if not text.endswith("\ufffd") and not (
-            len(v) == 1 and self._byte_decoder[v[0]] == 32
+            len(v) == 1 and self._byte_decoder.get(v[0]) == 32
         ):
             self.text += self._maybe_trim_space(text)
             self._unflushed = ""
@@ -423,9 +423,12 @@ def _is_bpe_decoder(decoder):
     return isinstance(decoder, dict) and decoder.get("type", None) == "ByteLevel"
 
 
-def load_tokenizer(
-    model_path, tokenizer_config_extra={}, return_tokenizer=True, eos_token_ids=None
-):
+def load(
+    model_path,
+    tokenizer_config_extra: Optional[Dict[str, Any]] = None,
+    return_tokenizer=True,
+    eos_token_ids=None,
+) -> TokenizerWrapper:
     """Load a huggingface tokenizer and try to infer the type of streaming
     detokenizer to use.
 
@@ -435,6 +438,7 @@ def load_tokenizer(
     detokenizer_class = NaiveStreamingDetokenizer
 
     tokenizer_file = model_path / "tokenizer.json"
+
     if tokenizer_file.exists():
         with open(tokenizer_file, "r", encoding="utf-8") as fid:
             try:
@@ -454,8 +458,9 @@ def load_tokenizer(
         eos_token_ids = [eos_token_ids]
 
     if return_tokenizer:
+        kwargs = tokenizer_config_extra or {}
         return TokenizerWrapper(
-            AutoTokenizer.from_pretrained(model_path, **tokenizer_config_extra),
+            AutoTokenizer.from_pretrained(model_path, **kwargs),
             detokenizer_class,
             eos_token_ids=eos_token_ids,
         )

mlx_lm/tuner/datasets.py

@@ -39,7 +39,7 @@ class TextDataset:
 class ChatDataset:
     """
     A dataset for chat data in the format of {"messages": [...]}
-    https://platform.openai.com/docs/guides/fine-tuning/example-format
+    https://platform.openai.com/docs/guides/supervised-fine-tuning#formatting-your-data
     """
 
     def __init__(
@@ -59,8 +59,14 @@ class ChatDataset:
         tools = d.get("tools", None)
         tokens = self.tokenizer.apply_chat_template(messages, tools=tools)
         if self.mask_prompt:
-            messages = messages[:-1]
-            offset = len(self.tokenizer.apply_chat_template(messages, tools=tools))
+            add_generation_prompt = messages[-1].get("role") == "assistant"
+            offset = len(
+                self.tokenizer.apply_chat_template(
+                    messages[:-1],
+                    tools=tools,
+                    add_generation_prompt=add_generation_prompt,
+                )
+            )
             return (tokens, offset)
         else:
             return (tokens, 0)
@@ -94,16 +100,16 @@ class CompletionsDataset:
         self.tokenizer = tokenizer
 
     def process(self, d):
-        tokens = self.tokenizer.apply_chat_template(
-            [
-                {"role": "user", "content": d[self.prompt_key]},
-                {"role": "assistant", "content": d[self.completion_key]},
-            ],
-        )
+        tools = d.get("tools", None)
+        messages = [
+            {"role": "user", "content": d[self.prompt_key]},
+            {"role": "assistant", "content": d[self.completion_key]},
+        ]
+        tokens = self.tokenizer.apply_chat_template(messages, tools=tools)
         if self.mask_prompt:
             offset = len(
                 self.tokenizer.apply_chat_template(
-                    [{"role": "user", "content": d[self.prompt_key]}]
+                    messages[0], tools=tools, add_generation_prompt=True
                 )
             )
             return (tokens, offset)

mlx_lm/tuner/dora.py

@@ -29,31 +29,29 @@ class DoRALinear(nn.Module):
         dora_lin.set_linear(linear)
         return dora_lin
 
-    def fuse(self, de_quantize: bool = False):
+    def fuse(self, dequantize: bool = False):
         linear = self.linear
         bias = "bias" in linear
         weight = self._dequantized_weight()
 
-        # Use the same type as the linear weight
-        dtype = weight.dtype
-
         output_dims, input_dims = weight.shape
         fused_linear = nn.Linear(input_dims, output_dims, bias=False)
 
-        lora_b = (self.scale * self.lora_b.T).astype(dtype)
-        lora_a = self.lora_a.T.astype(dtype)
-        weight = weight + lora_b @ lora_a
+        lora_b = self.scale * self.lora_b.T
+        lora_a = self.lora_a.T
+        weight = weight + (lora_b @ lora_a).astype(weight.dtype)
         norm_scale = self.m / mx.linalg.norm(weight, axis=1)
         fused_linear.weight = norm_scale[:, None] * weight
 
         if bias:
             fused_linear.bias = linear.bias
 
-        if self._is_quantized() and not de_quantize:
+        if self._is_quantized() and not dequantize:
             fused_linear = nn.QuantizedLinear.from_linear(
                 fused_linear,
-                linear.group_size,
-                linear.bits,
+                group_size=linear.group_size,
+                bits=linear.bits,
+                mode=linear.mode,
             )
         return fused_linear
 
@@ -101,8 +99,9 @@ class DoRALinear(nn.Module):
                 weight,
                 self.linear.scales,
                 self.linear.biases,
-                self.linear.group_size,
-                self.linear.bits,
+                group_size=self.linear.group_size,
+                bits=self.linear.bits,
+                mode=self.linear.mode,
             )
         return weight
 
@@ -151,19 +150,16 @@ class DoRAEmbedding(nn.Module):
         dora_embedding.set_embedding(embedding)
         return dora_embedding
 
-    def fuse(self, de_quantize: bool = False):
+    def fuse(self, dequantize: bool = False):
         embedding = self.embedding
         weight = embedding.weight
 
-        # Use the same type as the linear weight if not quantized
-        dtype = weight.dtype
-
         num_embeddings, dims = weight.shape
         fused_embedding = nn.Embedding(num_embeddings, dims)
 
-        lora_a = (self.scale * self.lora_a).astype(dtype)
-        lora_b = self.lora_b.astype(dtype)
-        weight = weight + lora_a @ lora_b
+        lora_a = self.scale * self.lora_a
+        lora_b = self.lora_b
+        weight = weight + (lora_a @ lora_b).astype(weight.dtype)
         norm_scale = self.m / mx.linalg.norm(weight, axis=1)
         fused_embedding.weight = norm_scale[:, None] * weight
 

mlx_lm/tuner/lora.py

@@ -31,37 +31,35 @@ class LoRALinear(nn.Module):
         lora_lin.linear = linear
         return lora_lin
 
-    def fuse(self, de_quantize: bool = False):
+    def fuse(self, dequantize: bool = False):
         linear = self.linear
         bias = "bias" in linear
         weight = linear.weight
         is_quantized = isinstance(linear, nn.QuantizedLinear)
 
-        # 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,
+                group_size=linear.group_size,
+                bits=linear.bits,
+                mode=linear.mode,
             )
         output_dims, input_dims = weight.shape
         fused_linear = nn.Linear(input_dims, output_dims, bias=bias)
 
-        delta = ((self.scale * self.lora_b.T) @ self.lora_a.T).astype(dtype)
+        delta = ((self.scale * self.lora_b.T) @ self.lora_a.T).astype(weight.dtype)
         fused_linear.weight = weight + delta
         if bias:
             fused_linear.bias = linear.bias
 
-        if is_quantized and not de_quantize:
+        if is_quantized and not dequantize:
             fused_linear = nn.QuantizedLinear.from_linear(
                 fused_linear,
                 linear.group_size,
                 linear.bits,
+                mode=linear.mode,
             )
 
         return fused_linear
@@ -119,35 +117,34 @@ class LoRASwitchLinear(nn.Module):
         lora_lin.linear = linear
         return lora_lin
 
-    def fuse(self, de_quantize: bool = False):
+    def fuse(self, dequantize: bool = False):
         linear = self.linear
         bias = "bias" in linear
         weight = linear.weight
         is_quantized = isinstance(linear, QuantizedSwitchLinear)
 
-        # Use the same type as the linear weight if not quantized
-        dtype = weight.dtype
-
         if is_quantized:
-            dtype = mx.float16
             weight = mx.dequantize(
                 weight,
                 linear.scales,
                 linear.biases,
-                linear.group_size,
-                linear.bits,
+                group_size=linear.group_size,
+                bits=linear.bits,
+                mode=linear.mode,
             )
         num_experts, output_dims, input_dims = weight.shape
         fused_linear = SwitchLinear(input_dims, output_dims, num_experts, bias=bias)
 
-        lora_b = (self.scale * self.lora_b).astype(dtype)
-        lora_a = self.lora_a.reshape(num_experts, -1, input_dims).astype(dtype)
-        fused_linear.weight = weight + lora_b @ lora_a
+        lora_b = self.scale * self.lora_b
+        lora_a = self.lora_a.reshape(num_experts, -1, input_dims)
+        fused_linear.weight = weight + (lora_b @ lora_a).astype(weight.dtype)
         if bias:
             fused_linear.bias = linear.bias
 
-        if is_quantized and not de_quantize:
-            fused_linear = fused_linear.to_quantized(linear.group_size, linear.bits)
+        if is_quantized and not dequantize:
+            fused_linear = fused_linear.to_quantized(
+                group_size=linear.group_size, bits=linear.bits, mode=linear.mode
+            )
 
         return fused_linear
 
@@ -176,19 +173,25 @@ class LoRASwitchLinear(nn.Module):
         self.lora_a = mx.random.uniform(
             low=-scale,
             high=scale,
-            shape=(r * num_experts, input_dims),
+            shape=(num_experts, r, input_dims),
         )
         self.lora_b = mx.zeros(shape=(num_experts, output_dims, r))
         self.num_experts = num_experts
 
-    def __call__(self, x, indices):
-        shape = x.shape[:-3] + (self.num_experts, -1)
-
-        y = self.linear(x, indices)
-        z = (self.dropout(x) @ self.lora_a.T).reshape(shape)
-        z = mx.take_along_axis(z, indices[..., None], axis=-2)
-        z = z[..., None, :] @ self.lora_b[indices].swapaxes(-2, -1)
-
+    def __call__(self, x, indices, sorted_indices=False):
+        y = self.linear(x, indices, sorted_indices=sorted_indices)
+        z = mx.gather_mm(
+            self.dropout(x),
+            self.lora_a.swapaxes(-1, -2),
+            rhs_indices=indices,
+            sorted_indices=sorted_indices,
+        )
+        z = mx.gather_mm(
+            z,
+            self.lora_b.swapaxes(-1, -2),
+            rhs_indices=indices,
+            sorted_indices=sorted_indices,
+        )
         return y + (self.scale * z).astype(x.dtype)
 
 
@@ -213,35 +216,33 @@ class LoRAEmbedding(nn.Module):
         lora_embedding.embedding = embedding
         return lora_embedding
 
-    def fuse(self, de_quantize: bool = False):
+    def fuse(self, dequantize: bool = False):
         embedding = self.embedding
         weight = embedding.weight
         is_quantized = isinstance(embedding, nn.QuantizedEmbedding)
 
-        # Use the same type as the linear weight if not quantized
-        dtype = weight.dtype
-
         if is_quantized:
-            dtype = embedding.scales.dtype
             weight = mx.dequantize(
                 weight,
                 embedding.scales,
                 embedding.biases,
-                embedding.group_size,
-                embedding.bits,
+                group_size=embedding.group_size,
+                bits=embedding.bits,
+                mode=embedding.mode,
             )
         num_embeddings, dims = weight.shape
         fused_embedding = nn.Embedding(num_embeddings, dims)
 
-        lora_a = (self.scale * self.lora_a).astype(dtype)
-        lora_b = self.lora_b.astype(dtype)
-        fused_embedding.weight = weight + lora_a @ lora_b
+        lora_a = self.scale * self.lora_a
+        lora_b = self.lora_b
+        fused_embedding.weight = weight + (lora_a @ lora_b).astype(weight.dtype)
 
-        if is_quantized and not de_quantize:
+        if is_quantized and not dequantize:
             fused_embedding = nn.QuantizedEmbedding.from_embedding(
                 fused_embedding,
-                embedding.group_size,
-                embedding.bits,
+                group_size=embedding.group_size,
+                bits=embedding.bits,
+                mode=embedding.mode,
             )
 
         return fused_embedding

mlx_lm/tuner/trainer.py

@@ -10,7 +10,7 @@ import mlx.core as mx
 import mlx.nn as nn
 import numpy as np
 from mlx.nn.utils import average_gradients
-from mlx.utils import tree_flatten
+from mlx.utils import tree_flatten, tree_map
 from tqdm import tqdm
 
 from .callbacks import TrainingCallback
@@ -64,6 +64,12 @@ class TrainingArgs:
         default=False,
         metadata={"help": "Use gradient checkpointing to reduce memory use."},
     )
+    grad_accumulation_steps: int = field(
+        default=1,
+        metadata={
+            "help": "Number of steps to accumulate gradients before applying an optimizer update."
+        },
+    )
 
 
 def default_loss(model, batch, lengths):
@@ -86,7 +92,9 @@ def iterate_batches(
     dataset,
     batch_size,
     max_seq_length,
-    train=False,
+    loop=False,
+    seed=None,
+    comm_group=None,
 ):
     # Sort by length:
     if isinstance(dataset, CacheDataset):
@@ -102,8 +110,12 @@ def iterate_batches(
 
     # If running in distributed mode (N machines) then each one should skip N-1
     # samples
-    offset = mx.distributed.init().rank()
-    step = mx.distributed.init().size()
+    if comm_group is not None:
+        offset = comm_group.rank()
+        step = comm_group.size()
+    else:
+        offset = 0
+        step = 1
     if batch_size % step != 0:
         raise ValueError("The batch size must be divisible by the number of workers")
 
@@ -112,7 +124,8 @@ def iterate_batches(
         idx[i + offset : i + offset + batch_size : step]
         for i in range(0, len(idx) - batch_size + 1, batch_size)
     ]
-
+    if seed:
+        np.random.seed(seed)
     while True:
         indices = np.random.permutation(len(batch_idx))
         for i in indices:
@@ -145,7 +158,7 @@ def iterate_batches(
             batch = mx.array(batch_arr)
             yield batch, mx.array(list(zip(offsets, lengths)))
 
-        if not train:
+        if not loop:
             break
 
 
@@ -171,6 +184,7 @@ def evaluate(
                 dataset=dataset,
                 batch_size=batch_size,
                 max_seq_length=max_seq_length,
+                comm_group=mx.distributed.init(),
             ),
         ),
         desc="Calculating loss...",
@@ -209,22 +223,29 @@ def train(
     if args.grad_checkpoint:
         grad_checkpoint(model.layers[0])
 
+    loss_value_and_grad = nn.value_and_grad(model, loss)
+
+    grad_accum_steps = args.grad_accumulation_steps
+    if grad_accum_steps < 1:
+        raise ValueError("grad_accumulation_steps must be at least 1")
+
     state = [model.state, optimizer.state, mx.random.state]
 
     @partial(mx.compile, inputs=state, outputs=state)
-    def step(batch):
-        # Forward and backward pass
+    def step(batch, prev_grad, do_update):
         (lvalue, toks), grad = loss_value_and_grad(model, *batch)
 
-        # All reduce the gradients if running in distributed mode
-        grad = average_gradients(grad)
+        if prev_grad is not None:
+            grad = tree_map(lambda x, y: x + y, grad, prev_grad)
 
-        # Model update
-        optimizer.update(model, grad)
+        if do_update:
+            grad = average_gradients(grad)
+            if grad_accum_steps > 1:
+                grad = tree_map(lambda x: x / grad_accum_steps, grad)
+            optimizer.update(model, grad)
+            grad = None
 
-        return lvalue, toks
-
-    loss_value_and_grad = nn.value_and_grad(model, loss)
+        return lvalue, toks, grad
 
     model.train()
     losses = 0
@@ -232,6 +253,8 @@ def train(
     steps = 0
     trained_tokens = 0
     train_time = 0
+    grad_accum = None
+
     # Main training loop
     for it, batch in zip(
         range(1, args.iters + 1),
@@ -239,7 +262,8 @@ def train(
             dataset=train_dataset,
             batch_size=args.batch_size,
             max_seq_length=args.max_seq_length,
-            train=True,
+            loop=True,
+            comm_group=world,
         ),
     ):
         tic = time.perf_counter()
@@ -276,11 +300,16 @@ def train(
 
             tic = time.perf_counter()
 
-        lvalue, toks = step(batch)
+        lvalue, toks, grad_accum = step(
+            batch,
+            grad_accum,
+            it % grad_accum_steps == 0,
+        )
+
         losses += lvalue
         n_tokens += toks
         steps += 1
-        mx.eval(state, losses, n_tokens)
+        mx.eval(state, losses, n_tokens, grad_accum)
         train_time += time.perf_counter() - tic
 
         # Report training loss if needed

mlx_lm/tuner/utils.py

@@ -10,6 +10,7 @@ import mlx.optimizers as opt
 from mlx.utils import tree_flatten, tree_unflatten
 
 from ..models.switch_layers import QuantizedSwitchLinear, SwitchLinear
+from ..utils import get_total_parameters
 from .dora import DoRAEmbedding, DoRALinear
 from .lora import LoRAEmbedding, LoRALinear, LoRASwitchLinear
 
@@ -81,140 +82,23 @@ def linear_to_lora_layers(
             dropout=config["dropout"],
         )
 
-    keys = config.get("keys", None)
-    if keys is not None:
-        keys = set(keys)
-    elif model.model_type in {
-        "mistral",
-        "mistral3",
-        "llama",
-        "llama4_text",
-        "lfm2",
-        "phi",
-        "mixtral",
-        "nemotron",
-        "stablelm",
-        "hunyuan",
-        "qwen2",
-        "qwen2_moe",
-        "qwen3",
-        "qwen3_moe",
-        "phimoe",
-        "gemma",
-        "gemma2",
-        "gemma3",
-        "gemma3_text",
-        "granite",
-        "helium",
-        "starcoder2",
-        "cohere",
-        "cohere2",
-        "minicpm",
-        "minicpm3",
-        "minicpm4",
-        "deepseek",
-        "olmo2",
-        "olmoe",
-        "internlm3",
-        "glm4",
-        "glm",
-        "mimo",
-        "ernie4_5",
-        "dots1",
-        "smollm3",
-        "exaone4",
-        "hunyuan_v1_dense",
-        "gpt_oss",
-        "ernie4_5_moe",
-        "granitemoe",
-        "longcat_flash",
-        "seed_oss",
-        "apertus",
-        "qwen3_next",
-        "Klear",
-        "lille-130m",
-    }:
-        keys = {"self_attn.q_proj", "self_attn.v_proj"}
-        if model.model_type in ["mixtral", "phimoe"]:
-            keys.add("block_sparse_moe.gate")
-        if model.model_type in ["qwen2_moe", "qwen3_next"]:
-            keys.add("mlp.gate")
-            keys.add("mlp.shared_expert_gate")
-        if model.model_type in ["olmoe", "qwen3_moe", "dots1", "Klear"]:
-            keys.add("mlp.gate")
-        if model.model_type in ["longcat_flash"]:
-            keys.add("mlp.router.classifier")
-        if model.model_type == "lille-130m":
-            keys.add("attention.qkv_proj")
-            keys.add("attention.out_proj")
-            keys.add("feed_forward.gate_proj")
-            keys.add("feed_forward.up_proj")
-            keys.add("feed_forward.down_proj")
-        elif model.model_type == "qwen3_next":
-            keys.add("linear_attn.in_proj_qkvz")
-            keys.add("linear_attn.out_proj")
-            keys.add("linear_attn.in_proj_ba")
-            keys.add("linear_attn.dt_bias")
-            keys.add("self_attn.q_proj")
-            keys.add("self_attn.k_proj")
-            keys.add("self_attn.v_proj")
-            keys.add("self_attn.o_proj")
-    elif model.model_type == "gpt_bigcode":
-        keys = {"attn.c_attn"}
-    elif model.model_type == "gpt2":
-        keys = {"attn.c_attn"}
-    elif model.model_type == "gpt_neox":
-        keys = {"attention.query_key_value"}
-    elif model.model_type == "olmo":
-        keys = {"att_proj"}
-    elif model.model_type == "openelm":
-        keys = {"attn.qkv_proj"}
-    elif model.model_type == "phi3":
-        keys = {"self_attn.qkv_proj"}
-    elif model.model_type == "phi-msft":
-        keys = {"mixer.Wqkv", "moe.gate"}
-    elif model.model_type == "dbrx":
-        keys = {"norm_attn_norm.attn.Wqkv", "ffn.router.layer"}
-    elif model.model_type == "internlm2":
-        keys = {"attention.wqkv", "attention.wo"}
-    elif model.model_type in {
-        "deepseek_v2",
-        "deepseek_v3",
-        "longcat_flash",
-        "minicpm3",
-    }:
-        keys = {
-            "self_attn.q_proj",
-            "self_attn.q_a_proj",
-            "self_attn.q_b_proj",
-            "self_attn.kv_a_proj_with_mqa",
-            "self_attn.kv_b_proj",
-        }
-    elif model.model_type == "mamba":
-        keys = {
-            "mixer.in_proj",
-            "mixer.x_proj",
-            "mixer.dt_proj",
-            "mixer.out_proj",
-        }
-    elif model.model_type == "mamba2":
-        keys = {
-            "mixer.in_proj",
-            "mixer.out_proj",
-        }
-    elif model.model_type == "exaone":
-        keys = {"attn.attention.q_proj", "attn.attention.v_proj"}
-    elif model.model_type == "bailing_moe":
-        keys = {"attention.query_key_value", "attention.dense"}
-    elif model.model_type == "nemotron_h":
-        keys.add("mixer.in_proj")
-        keys.add("mixer.out_proj")
-        keys.add("mixer.q_proj")
-        keys.add("mixer.k_proj")
-        keys.add("mixer.v_proj")
-        keys.add("mixer.o_proj")
-    else:
-        raise ValueError(f"Lora does not support {model.model_type}")
+    if (keys := config.get("keys", None)) is None:
+        keys = set()
+
+        def get_keys_for_lora(p, m):
+            types = (
+                nn.Linear,
+                nn.QuantizedLinear,
+                SwitchLinear,
+                QuantizedSwitchLinear,
+                nn.Embedding,
+                nn.QuantizedEmbedding,
+            )
+            if hasattr(m, "to_lora") or isinstance(m, types):
+                keys.add(p)
+
+        for l in model.layers:
+            l.apply_to_modules(get_keys_for_lora)
 
     for l in model.layers[-max(num_layers, 0) :]:
         lora_layers = [(k, to_lora(m)) for k, m in l.named_modules() if k in keys]
@@ -254,49 +138,6 @@ def load_adapters(model: nn.Module, adapter_path: str) -> nn.Module:
     return model
 
 
-def dequantize(model: nn.Module) -> nn.Module:
-    """
-    Dequantize the quantized linear layers in the model.
-
-    Args:
-        model (nn.Module): The model with quantized linear layers.
-
-    Returns:
-        nn.Module: The model with dequantized layers.
-    """
-    dequantize_layers = []
-    for name, module in model.named_modules():
-        bias = "bias" in module
-        if isinstance(module, nn.QuantizedLinear):
-            cls = nn.Linear
-            kwargs = {"bias": bias}
-        elif isinstance(module, nn.QuantizedEmbedding):
-            kwargs = {}
-            cls = nn.Embedding
-        elif isinstance(module, QuantizedSwitchLinear):
-            kwargs = {"bias": bias}
-            cls = SwitchLinear
-        else:
-            continue
-        weight = mx.dequantize(
-            module.weight,
-            module.scales,
-            module.biases,
-            module.group_size,
-            module.bits,
-        )
-        args = weight.shape[::-1]
-        m = cls(*args, **kwargs)
-        if bias:
-            m.bias = module.bias
-        m.weight = weight
-        dequantize_layers.append((name, m))
-
-    if len(dequantize_layers) > 0:
-        model.update_modules(tree_unflatten(dequantize_layers))
-    return model
-
-
 def remove_lora_layers(model: nn.Module) -> nn.Module:
     """
     Remove the LoRA layers from the model.
@@ -316,20 +157,6 @@ def remove_lora_layers(model: nn.Module) -> nn.Module:
     return model
 
 
-def get_total_parameters(model):
-    leaf_modules = tree_flatten(
-        model.leaf_modules(), is_leaf=lambda m: isinstance(m, nn.Module)
-    )
-
-    def nparams(m):
-        if hasattr(m, "bits"):
-            n = 0 if not hasattr(m, "bias") else m.bias.size
-            return n + m.weight.size * 32 // m.bits
-        return sum(v.size for _, v in tree_flatten(m.parameters()))
-
-    return sum(nparams(m) for _, m in leaf_modules)
-
-
 def print_trainable_parameters(model):
     total_p = get_total_parameters(model) / 1e6
     trainable_p = (

mlx_lm/utils.py

@@ -7,6 +7,7 @@ import inspect
 import json
 import logging
 import os
+import resource
 import shutil
 from pathlib import Path
 from textwrap import dedent
@@ -14,6 +15,7 @@ from typing import (
     Any,
     Callable,
     Dict,
+    List,
     Optional,
     Tuple,
     Type,
@@ -31,17 +33,19 @@ if os.getenv("MLXLM_USE_MODELSCOPE", "False").lower() == "true":
 else:
     from huggingface_hub import snapshot_download
 
-from mlx.utils import tree_flatten, tree_map, tree_reduce
-from transformers import PreTrainedTokenizer
+# For large models with lots of files
+resource.setrlimit(resource.RLIMIT_NOFILE, (2048, 4096))
+
+from mlx.utils import tree_flatten, tree_map, tree_reduce, tree_unflatten
 
 # Local imports
-from .tokenizer_utils import TokenizerWrapper, load_tokenizer
-from .tuner.utils import dequantize as dequantize_model
-from .tuner.utils import get_total_parameters, load_adapters
+from .tokenizer_utils import TokenizerWrapper
+from .tokenizer_utils import load as _load_tokenizer
 
 # Constants
 MODEL_REMAPPING = {
     "mistral": "llama",
+    "llava": "mistral3",
     "phi-msft": "phixtral",
     "falcon_mamba": "mamba",
     "kimi_k2": "deepseek_v3",
@@ -73,6 +77,20 @@ def _get_classes(config: dict):
     return arch.Model, arch.ModelArgs
 
 
+def get_total_parameters(model):
+    leaf_modules = tree_flatten(
+        model.leaf_modules(), is_leaf=lambda m: isinstance(m, nn.Module)
+    )
+
+    def nparams(m):
+        if hasattr(m, "bits"):
+            n = 0 if not hasattr(m, "bias") else m.bias.size
+            return n + m.weight.size * 32 // m.bits
+        return sum(v.size for _, v in tree_flatten(m.parameters()))
+
+    return sum(nparams(m) for _, m in leaf_modules)
+
+
 def compute_bits_per_weight(model):
     model_bytes = tree_reduce(
         lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, model, 0
@@ -81,9 +99,11 @@ def compute_bits_per_weight(model):
     return model_bytes * 8 / model_params
 
 
-def get_model_path(
-    path_or_hf_repo: str, revision: Optional[str] = None
-) -> Tuple[Path, Optional[str]]:
+def _download(
+    path_or_hf_repo: str,
+    revision: Optional[str] = None,
+    allow_patterns: List[str] = None,
+) -> Path:
     """
     Ensures the model is available locally. If the path does not exist locally,
     it is downloaded from the Hugging Face Hub.
@@ -93,39 +113,35 @@ def get_model_path(
         revision (str, optional): A revision id which can be a branch name, a tag, or a commit hash.
 
     Returns:
-        Tuple[Path, str]: A tuple containing the local file path and the Hugging Face repo ID.
+        Path: The local file path.
     """
     model_path = Path(path_or_hf_repo)
 
     if not model_path.exists():
-        hf_path = path_or_hf_repo
+        allow_patterns = allow_patterns or [
+            "*.json",
+            "model*.safetensors",
+            "*.py",
+            "tokenizer.model",
+            "*.tiktoken",
+            "tiktoken.model",
+            "*.txt",
+            "*.jsonl",
+            "*.jinja",
+        ]
         model_path = Path(
             snapshot_download(
                 path_or_hf_repo,
                 revision=revision,
-                allow_patterns=[
-                    "*.json",
-                    "model*.safetensors",
-                    "*.py",
-                    "tokenizer.model",
-                    "*.tiktoken",
-                    "tiktoken.model",
-                    "*.txt",
-                    "*.jsonl",
-                    "*.jinja",
-                ],
+                allow_patterns=allow_patterns,
             )
         )
-    else:
-        from huggingface_hub import ModelCard
 
-        card_path = model_path / "README.md"
-        if card_path.is_file():
-            card = ModelCard.load(card_path)
-            hf_path = card.data.base_model
-        else:
-            hf_path = None
-    return model_path, hf_path
+    return model_path
+
+
+def hf_repo_to_path(hf_repo):
+    return Path(snapshot_download(hf_repo, local_files_only=True))
 
 
 def load_config(model_path: Path) -> dict:
@@ -142,7 +158,7 @@ def load_model(
     model_path: Path,
     lazy: bool = False,
     strict: bool = True,
-    model_config: dict = {},
+    model_config: Optional[Dict[str, Any]] = None,
     get_model_classes: Callable[[dict], Tuple[Type[nn.Module], Type]] = _get_classes,
 ) -> Tuple[nn.Module, dict]:
     """
@@ -169,7 +185,8 @@ def load_model(
         ValueError: If the model class or args class are not found or cannot be instantiated.
     """
     config = load_config(model_path)
-    config.update(model_config)
+    if model_config is not None:
+        config.update(model_config)
 
     weight_files = glob.glob(str(model_path / "model*.safetensors"))
 
@@ -221,6 +238,11 @@ def load_model(
             config["quantization"] = quantization
             config["quantization_config"] = quantization
             _quantize(quantization)
+        elif quant_method == "compressed-tensors":
+            quantization = {"group_size": 32, "bits": 4, "mode": "affine"}
+            config["quantization"] = quantization
+            config["quantization_config"] = quantization
+            _quantize(quantization)
 
     model.load_weights(list(weights.items()), strict=strict)
 
@@ -231,13 +253,46 @@ def load_model(
     return model, config
 
 
+def load_adapters(model: nn.Module, adapter_path: str) -> nn.Module:
+    from .tuner.utils import load_adapters as _load_adapters
+
+    return _load_adapters(model, adapter_path)
+
+
+def load_tokenizer(model_path, tokenizer_config_extra=None, eos_token_ids=None):
+    """Load a huggingface tokenizer and try to infer the type of streaming
+    detokenizer to use.
+    """
+    model_path = _download(
+        model_path,
+        allow_patterns=[
+            "*.json",
+            "*.py",
+            "tokenizer.model",
+            "*.tiktoken",
+            "tiktoken.model",
+            "*.txt",
+            "*.jsonl",
+            "*.jinja",
+        ],
+    )
+    return _load_tokenizer(
+        model_path, tokenizer_config_extra, eos_token_ids=eos_token_ids
+    )
+
+
 def load(
     path_or_hf_repo: str,
-    tokenizer_config={},
-    model_config={},
+    tokenizer_config: Optional[Dict[str, Any]] = None,
+    model_config: Optional[Dict[str, Any]] = None,
     adapter_path: Optional[str] = None,
     lazy: bool = False,
-) -> Tuple[nn.Module, TokenizerWrapper]:
+    return_config: bool = False,
+    revision: Optional[str] = None,
+) -> Union[
+    Tuple[nn.Module, TokenizerWrapper],
+    Tuple[nn.Module, TokenizerWrapper, Dict[str, Any]],
+]:
     """
     Load the model and tokenizer from a given path or a huggingface repository.
 
@@ -252,16 +307,19 @@ def load(
         lazy (bool): If ``False`` eval the model parameters to make sure they are
             loaded in memory before returning, otherwise they will be loaded
             when needed. Default: ``False``
+        return_config (bool: If ``True`` return the model config as the last item..
+        revision (str, optional): A revision id which can be a branch name, a tag, or a commit hash.
     Returns:
-        Tuple[nn.Module, TokenizerWrapper]: A tuple containing the loaded model and tokenizer.
+        Union[Tuple[nn.Module, TokenizerWrapper], Tuple[nn.Module, TokenizerWrapper, Dict[str, Any]]]:
+            A tuple containing the loaded model, tokenizer and, if requested, the model config.
 
     Raises:
         FileNotFoundError: If config file or safetensors are not found.
         ValueError: If model class or args class are not found.
     """
-    model_path, _ = get_model_path(path_or_hf_repo)
+    model_path = _download(path_or_hf_repo, revision=revision)
 
-    model, config = load_model(model_path, lazy)
+    model, config = load_model(model_path, lazy, model_config=model_config)
     if adapter_path is not None:
         model = load_adapters(model, adapter_path)
         model.eval()
@@ -269,19 +327,66 @@ def load(
         model_path, tokenizer_config, eos_token_ids=config.get("eos_token_id", None)
     )
 
-    return model, tokenizer
+    if return_config:
+        return model, tokenizer, config
+    else:
+        return model, tokenizer
 
 
-def fetch_from_hub(
-    model_path: Path, lazy: bool = False, trust_remote_code: bool = False
-) -> Tuple[nn.Module, dict, PreTrainedTokenizer]:
-    model, config = load_model(model_path, lazy)
+def pipeline_load(repo, return_config=False):
+    # Get model path with everything but weight safetensors
+    model_path = _download(
+        repo,
+        allow_patterns=[
+            "*.json",
+            "*.py",
+            "tokenizer.model",
+            "*.tiktoken",
+            "tiktoken.model",
+            "*.txt",
+            "*.jsonl",
+            "*.jinja",
+        ],
+    )
+
+    # Lazy load and shard model to figure out which weights we need
+    model, config = load_model(model_path, lazy=True, strict=False)
+
+    group = mx.distributed.init()
+    rank = group.rank()
+    model.model.pipeline(group)
+
+    # Figure out which files we need for the local shard
+    with open(model_path / "model.safetensors.index.json", "r") as fid:
+        weight_index = json.load(fid)["weight_map"]
+
+    local_files = set()
+    for k, _ in tree_flatten(model.parameters()):
+        if file_name := weight_index.get(k, None) is None:
+            raise ValueError(
+                "Pipeline loading is only supported for MLX converted models."
+            )
+        local_files.add(weight_index[k])
+
+    # Download weights for local shard
+    _download(repo, allow_patterns=local_files)
+
+    # Load and shard the model, and load the weights
     tokenizer = load_tokenizer(
         model_path,
+        {"trust_remote_code": True},
         eos_token_ids=config.get("eos_token_id", None),
-        tokenizer_config_extra={"trust_remote_code": trust_remote_code},
     )
-    return model, config, tokenizer
+    model, _ = load_model(model_path, lazy=True, strict=False)
+    model.model.pipeline(group)
+    mx.eval(model.parameters())
+
+    # Synchronize processes to avoid timeout
+    mx.eval(mx.distributed.all_sum(mx.array(1.0), stream=mx.cpu))
+    if return_config:
+        return model, tokenizer, config
+    else:
+        return model, tokenizer
 
 
 def make_shards(weights: dict, max_file_size_gb: int = MAX_FILE_SIZE_GB) -> list:
@@ -308,24 +413,28 @@ def make_shards(weights: dict, max_file_size_gb: int = MAX_FILE_SIZE_GB) -> list
     return shards
 
 
-def create_model_card(path: Union[str, Path], hf_path: Union[str, Path]):
+def create_model_card(path: Union[str, Path], hf_path: Union[str, Path, None]):
     """
     Uploads the model to Hugging Face hub.
 
     Args:
         path (Union[str, Path]): Local path to the model.
-        hf_path (Union[str, Path]): Path to the original Hugging Face model.
+        hf_path (Union[str, Path, None]): Path to the original Hugging Face model.
     """
-    from huggingface_hub import ModelCard
+    from huggingface_hub import ModelCard, ModelCardData
 
-    card = ModelCard.load(hf_path)
+    if hf_path is None:
+        card = ModelCard.from_template(ModelCardData(language="en"))
+    else:
+        card = ModelCard.load(hf_path)
     card.data.library_name = "mlx"
     card.data.pipeline_tag = "text-generation"
     if card.data.tags is None:
         card.data.tags = ["mlx"]
     elif "mlx" not in card.data.tags:
         card.data.tags += ["mlx"]
-    card.data.base_model = str(hf_path)
+    if hf_path is not None:
+        card.data.base_model = str(hf_path)
     card.text = ""
     card.save(os.path.join(path, "README.md"))
 
@@ -345,15 +454,22 @@ def upload_to_hub(path: str, upload_repo: str):
     logging.set_verbosity_info()
     card_path = Path(path) / "README.md"
     card = ModelCard.load(card_path)
-    hf_path = card.data.base_model
-    card.text = dedent(
-        f"""
-        # {upload_repo}
 
+    hf_path = card.data.base_model
+
+    if hf_path is not None:
+        provenance = f"""
         This model [{upload_repo}](https://huggingface.co/{upload_repo}) was
         converted to MLX format from [{hf_path}](https://huggingface.co/{hf_path})
         using mlx-lm version **{__version__}**.
+        """
+    else:
+        provenance = ""
 
+    card.text = dedent(
+        f"""
+        # {upload_repo}
+        {provenance}
         ## Use with mlx
 
         ```bash
@@ -455,9 +571,7 @@ def quantize_model(
     group_size: int,
     bits: int,
     mode: str = "affine",
-    quant_predicate: Optional[
-        Callable[[str, nn.Module, dict], Union[bool, dict]]
-    ] = None,
+    quant_predicate: Optional[Callable[[str, nn.Module], Union[bool, dict]]] = None,
 ) -> Tuple[nn.Module, dict]:
     """
     Applies quantization to the model weights.
@@ -468,11 +582,10 @@ def quantize_model(
         group_size (int): Group size for quantization.
         bits (int): Bits per weight for quantization.
         mode (str): The quantization mode.
-        quant_predicate (Callable): A callable that decides how
-            to quantize each layer based on the path.
-            Accepts the layer `path`, the `module` and the model `config`.
-            Returns either a bool to signify quantize/no quantize or
-            a dict of quantization parameters to pass to `to_quantized`.
+        quant_predicate (Callable): A callable that decides how to quantize
+          each layer based on the path. Accepts the layer `path` and the
+          `module`. Returns either a bool to signify quantize/no quantize or
+          a dict of quantization parameters to pass to `to_quantized`.
 
     Returns:
         Tuple: Tuple containing quantized model and config.
@@ -496,7 +609,7 @@ def quantize_model(
             return False
         bool_or_params = True
         if quant_predicate is not None:
-            bool_or_params = quant_predicate(path, module, config)
+            bool_or_params = quant_predicate(path, module)
         if isinstance(bool_or_params, dict):
             quantized_config["quantization"][path] = bool_or_params
         elif fine_grained_config and bool_or_params:
@@ -519,6 +632,52 @@ def quantize_model(
     return model, quantized_config
 
 
+def dequantize_model(model: nn.Module) -> nn.Module:
+    """
+    Dequantize the quantized layers in the model.
+
+    Args:
+        model (nn.Module): The model with quantized layers.
+
+    Returns:
+        nn.Module: The model with dequantized layers.
+    """
+    from .models.switch_layers import QuantizedSwitchLinear, SwitchLinear
+
+    dequantize_layers = []
+    for name, module in model.named_modules():
+        bias = "bias" in module
+        if isinstance(module, nn.QuantizedLinear):
+            cls = nn.Linear
+            kwargs = {"bias": bias}
+        elif isinstance(module, nn.QuantizedEmbedding):
+            kwargs = {}
+            cls = nn.Embedding
+        elif isinstance(module, QuantizedSwitchLinear):
+            kwargs = {"bias": bias}
+            cls = SwitchLinear
+        else:
+            continue
+        weight = mx.dequantize(
+            module.weight,
+            module.scales,
+            module.biases,
+            module.group_size,
+            module.bits,
+            module.mode,
+        )
+        args = weight.shape[::-1]
+        m = cls(*args, **kwargs)
+        if bias:
+            m.bias = module.bias
+        m.weight = weight
+        dequantize_layers.append((name, m))
+
+    if len(dequantize_layers) > 0:
+        model.update_modules(tree_unflatten(dequantize_layers))
+    return model
+
+
 def save_config(
     config: dict,
     config_path: Union[str, Path],
@@ -547,14 +706,20 @@ def save_config(
 
 def save(
     dst_path: Union[str, Path],
-    src_path: Union[str, Path],
+    src_path_or_repo: Union[str, Path],
     model: nn.Module,
     tokenizer: TokenizerWrapper,
     config: Dict[str, Any],
-    hf_repo: Optional[str] = None,
     donate_model: bool = True,
 ):
-    src_path = Path(src_path)
+
+    src_path = Path(src_path_or_repo)
+    if not src_path.exists():
+        hf_repo = src_path_or_repo
+        src_path = hf_repo_to_path(hf_repo)
+    else:
+        hf_repo = None
+
     dst_path = Path(dst_path)
     save_model(dst_path, model, donate_model=True)
     save_config(config, config_path=dst_path / "config.json")
@@ -564,8 +729,7 @@ def save(
         for file in glob.glob(str(src_path / p)):
             shutil.copy(file, dst_path)
 
-    if hf_repo is not None:
-        create_model_card(dst_path, hf_repo)
+    create_model_card(dst_path, hf_repo)
 
 
 def common_prefix_len(list1, list2):

requirements.txt

@@ -1,6 +0,0 @@
-mlx>=0.29.1
-numpy
-transformers>=4.39.3
-protobuf
-pyyaml
-jinja2

setup.py

@@ -6,12 +6,12 @@ from pathlib import Path
 from setuptools import setup
 
 package_dir = Path(__file__).parent / "mlx_lm"
-with open("requirements.txt") as fid:
-    requirements = [l.strip() for l in fid.readlines()]
-
 sys.path.append(str(package_dir))
+
 from _version import __version__
 
+MIN_MLX_VERSION = "0.29.2"
+
 setup(
     name="mlx-lm",
     version=__version__,
@@ -23,13 +23,23 @@ setup(
     author="MLX Contributors",
     url="https://github.com/ml-explore/mlx-lm",
     license="MIT",
-    install_requires=requirements,
+    install_requires=[
+        f"mlx>={MIN_MLX_VERSION}; platform_system == 'Darwin'",
+        "numpy",
+        "transformers>=4.39.3",
+        "sentencepiece",
+        "protobuf",
+        "pyyaml",
+        "jinja2",
+    ],
     packages=["mlx_lm", "mlx_lm.models", "mlx_lm.quant", "mlx_lm.tuner"],
     python_requires=">=3.8",
     extras_require={
         "test": ["datasets", "lm-eval"],
         "train": ["datasets", "tqdm"],
         "evaluate": ["lm-eval", "tqdm"],
+        "cuda": [f"mlx[cuda]>={MIN_MLX_VERSION}"],
+        "cpu": [f"mlx[cpu]>={MIN_MLX_VERSION}"],
     },
     entry_points={
         "console_scripts": [

tests/test_finetune.py

@@ -5,7 +5,7 @@ import sys
 import unittest
 from contextlib import contextmanager
 from io import StringIO
-from unittest.mock import MagicMock
+from unittest.mock import ANY, MagicMock
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -28,13 +28,6 @@ def swapped_with_identity(obj, func):
 
 
 class TestLora(unittest.TestCase):
-    def setUp(self):
-        self.capturedOutput = StringIO()
-        sys.stdout = self.capturedOutput
-
-    def tearDown(self):
-        sys.stdout = sys.__stdout__
-
     def test_llama(self):
         from mlx_lm.models import llama
 
@@ -48,7 +41,6 @@ class TestLora(unittest.TestCase):
             vocab_size=10_000,
             tie_word_embeddings=False,
         )
-
         lora_layers = 4
 
         def check_config(params, expected_trainable_parameters=None):
@@ -68,10 +60,13 @@ class TestLora(unittest.TestCase):
             self.assertEqual(trainable_params, expected_trainable_parameters)
 
         params = {"rank": 8, "dropout": 0.0, "scale": 10.0}
-        check_config(params)
+        nparams = (
+            args.hidden_size * 2 * 4 + (args.intermediate_size + args.hidden_size) * 3
+        ) * lora_layers
+        check_config(params, expected_trainable_parameters=nparams * params["rank"])
 
         params["rank"] = 1
-        check_config(params)
+        check_config(params, expected_trainable_parameters=nparams * params["rank"])
 
         params["keys"] = ["self_attn.k_proj"]
         check_config(params)
@@ -128,7 +123,7 @@ class TestLora(unittest.TestCase):
             embedding.bits,
         )
         lora_emb = LoRAEmbedding.from_base(embedding, r=8, dropout=0, scale=10)
-        new_embedding = lora_emb.fuse(de_quantize=True)
+        new_embedding = lora_emb.fuse(dequantize=True)
         self.assertTrue(mx.array_equal(dequantized_weight, new_embedding.weight))
         self.assertTrue(mx.array_equal(embedding(tokens), lora_emb(tokens)))
 
@@ -142,7 +137,7 @@ class TestLora(unittest.TestCase):
 
         # change the value of lora_b and the embeddings will no longer be equal
         lora_emb.lora_b = mx.random.uniform(shape=lora_emb.lora_b.shape)
-        new_embedding = lora_emb.fuse(de_quantize=True)
+        new_embedding = lora_emb.fuse(dequantize=True)
         self.assertFalse(mx.array_equal(dequantized_weight, new_embedding.weight))
         self.assertFalse(mx.array_equal(embedding(tokens), lora_emb(tokens)))
 
@@ -191,7 +186,7 @@ class TestDora(unittest.TestCase):
 
         dora_layers = 4
 
-        def check_config(params):
+        def check_config(params, expected_params=None):
             n_keys = 2
             if "keys" in params:
                 n_keys = len(params["keys"])
@@ -201,17 +196,29 @@ class TestDora(unittest.TestCase):
             trainable_params = sum(
                 v.size for _, v in tree_flatten(model.trainable_parameters())
             )
-            self.assertEqual(
-                trainable_params,
+            expected_params = expected_params or (
                 dora_layers
-                * (params["rank"] * hidden_size * 2 * n_keys + n_keys * hidden_size),
+                * (params["rank"] * hidden_size * 2 * n_keys + n_keys * hidden_size)
             )
+            self.assertEqual(trainable_params, expected_params)
 
         params = {"rank": 8, "alpha": 16, "dropout": 0.0, "scale": 10.0}
-        check_config(params)
+        nparams = (
+            args.hidden_size * 2 * 4 + (args.intermediate_size + args.hidden_size) * 3
+        )
+        nparams = (
+            nparams * params["rank"] + 5 * args.hidden_size + 2 * args.intermediate_size
+        ) * dora_layers
+        check_config(params, expected_params=nparams)
 
         params["rank"] = 1
-        check_config(params)
+        nparams = (
+            args.hidden_size * 2 * 4 + (args.intermediate_size + args.hidden_size) * 3
+        )
+        nparams = (
+            nparams * params["rank"] + 5 * args.hidden_size + 2 * args.intermediate_size
+        ) * dora_layers
+        check_config(params, expected_params=nparams * params["rank"])
 
         params["keys"] = ["self_attn.k_proj"]
         check_config(params)
@@ -293,7 +300,7 @@ class TestDora(unittest.TestCase):
         quantized_linear = nn.QuantizedLinear(in_dims, out_dims, bias=True)
         dora_quantized_linear = DoRALinear.from_base(quantized_linear, r)
         # Dequantize
-        to_linear_from_quantized = dora_quantized_linear.fuse(de_quantize=True)
+        to_linear_from_quantized = dora_quantized_linear.fuse(dequantize=True)
         self.assertTrue(
             mx.allclose(quantized_linear.bias, to_linear_from_quantized.bias)
         )
@@ -398,6 +405,7 @@ class TestScheduleConfig(unittest.TestCase):
             dataset=mock_dataset,
             batch_size=2,
             max_seq_length=2048,
+            comm_group=ANY,
         )
         self.assertEqual(mock_default_loss.call_count, 2)
 
@@ -434,6 +442,7 @@ class TestScheduleConfig(unittest.TestCase):
             dataset=mock_dataset,
             batch_size=2,
             max_seq_length=2048,
+            comm_group=ANY,
         )
         self.assertEqual(mock_default_loss.call_count, 3)
 

tests/test_generate.py

@@ -11,6 +11,7 @@ from mlx_lm.generate import (
     generate,
     stream_generate,
 )
+from mlx_lm.models.cache import RotatingKVCache
 from mlx_lm.sample_utils import make_logits_processors, make_sampler
 from mlx_lm.utils import load
 
@@ -301,6 +302,136 @@ class TestGenerate(unittest.TestCase):
             batch_tokens = batch_responses[uids[e]]
             self.assertEqual(tokens, batch_tokens)
 
+    def test_batch_sliding_window(self):
+        prompts = [
+            "Write a story about Einstein",
+            "Hi",
+            "What time is it?",
+            "How tall is Mt Everest?",
+        ]
+        prompts = [
+            self.tokenizer.apply_chat_template(
+                [{"role": "user", "content": p}],
+                tokenize=True,
+                add_generation_prompt=True,
+            )
+            for p in prompts
+        ]
+
+        self.model.make_cache = lambda: [
+            RotatingKVCache(max_size=4) for _ in self.model.layers
+        ]
+        batch_gen = BatchGenerator(
+            self.model,
+            stop_tokens=self.tokenizer.eos_token_ids,
+            max_tokens=10,
+            prefill_batch_size=1,
+            prefill_step_size=8,
+            completion_batch_size=2,
+        )
+        uids = batch_gen.insert(prompts)
+        batch_responses = {uid: [] for uid in uids}
+        while responses := batch_gen.next():
+            for r in responses:
+                batch_responses[r.uid].append(r.logprobs)
+
+        for e, uid in enumerate(uids):
+            for i, response in enumerate(
+                stream_generate(
+                    self.model,
+                    self.tokenizer,
+                    prompts[e],
+                    max_tokens=10,
+                )
+            ):
+                batch_logprobs = batch_responses[uid][i]
+                logprobs = response.logprobs
+                self.assertTrue(
+                    mx.allclose(batch_logprobs, logprobs, rtol=1e-4, atol=1e-4)
+                )
+
+        del self.model.make_cache
+
+    def test_batch_continued_generation(self):
+        for rotating in [False, True]:
+            if rotating:
+                self.model.make_cache = lambda: [
+                    RotatingKVCache(max_size=4) for _ in self.model.layers
+                ]
+
+            # Make the prompts
+            prompts_a = [
+                "Write a story about Einstein",
+                "Hi",
+                "What time is it?",
+                "How tall is Mt Everest?",
+            ]
+            prompts_a = [
+                self.tokenizer.apply_chat_template(
+                    [{"role": "user", "content": p}],
+                    tokenize=True,
+                    add_generation_prompt=True,
+                )
+                for p in prompts_a
+            ]
+            prompts_b = [
+                "Another one",
+                "sup?",
+                "And how about the date?",
+                "Mt Olympus?",
+            ]
+            prompts_b = [
+                self.tokenizer.apply_chat_template(
+                    [{"role": "user", "content": p}],
+                    tokenize=True,
+                    add_generation_prompt=True,
+                )
+                for p in prompts_b
+            ]
+
+            # Generate once
+            batch_gen = BatchGenerator(
+                self.model,
+                stop_tokens=self.tokenizer.eos_token_ids,
+                max_tokens=10,
+                prefill_batch_size=1,
+                prefill_step_size=8,
+                completion_batch_size=2,
+            )
+            uids = batch_gen.insert(prompts_a)
+            caches = {uid: None for uid in uids}
+            while responses := batch_gen.next():
+                for r in responses:
+                    if r.finish_reason is not None:
+                        caches[r.uid] = r.prompt_cache
+            caches = [caches[uid] for uid in uids]
+
+            # Generate the 2nd time
+            uids = batch_gen.insert(prompts_b, caches=caches)
+            batch_responses = {uid: [] for uid in uids}
+            while responses := batch_gen.next():
+                for r in responses:
+                    batch_responses[r.uid].append(r.logprobs)
+
+            for e, uid in enumerate(uids):
+                for i, response in enumerate(
+                    stream_generate(
+                        self.model,
+                        self.tokenizer,
+                        prompts_b[e],
+                        max_tokens=10,
+                        prompt_cache=caches[e],
+                    )
+                ):
+                    batch_logprobs = batch_responses[uid][i]
+                    logprobs = response.logprobs
+                    self.assertTrue(
+                        mx.allclose(batch_logprobs, logprobs, rtol=1e-4, atol=1e-4)
+                    )
+
+            if rotating:
+                del self.model.make_cache
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_losses.py

@@ -11,6 +11,7 @@ class TestLosses(unittest.TestCase):
 
     def test_kl_div_loss(self):
         self.assertTrue(can_run_metal())
+        mx.random.seed(0)
 
         logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
         logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
@@ -23,6 +24,7 @@ class TestLosses(unittest.TestCase):
 
     def test_js_div_loss(self):
         self.assertTrue(can_run_metal())
+        mx.random.seed(0)
 
         logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
         logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
@@ -35,6 +37,7 @@ class TestLosses(unittest.TestCase):
 
     def test_kl_div_loss_vjp(self):
         self.assertTrue(can_run_metal())
+        mx.random.seed(0)
 
         logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
         logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
@@ -48,6 +51,7 @@ class TestLosses(unittest.TestCase):
 
     def test_js_div_loss_vjp(self):
         self.assertTrue(can_run_metal())
+        mx.random.seed(0)
 
         logits_q = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)
         logits_p = mx.random.uniform(shape=(4, 8, 4000), dtype=mx.float32)

tests/test_models.py

@@ -37,7 +37,7 @@ class TestModels(unittest.TestCase):
 
     def test_rotating_kv_cache(self):
         b, h, d = 1, 2, 32
-        cache = RotatingKVCache(max_size=8, step=4)
+        cache = RotatingKVCache(max_size=8)
 
         k = mx.random.uniform(shape=(b, h, 2, d))
         v = mx.random.uniform(shape=(b, h, 2, d))
@@ -70,7 +70,7 @@ class TestModels(unittest.TestCase):
             idx %= 8
 
         # Try with nonzero keep
-        cache = RotatingKVCache(max_size=8, step=4, keep=2)
+        cache = RotatingKVCache(max_size=8, keep=2)
 
         # Check a large update
         k = mx.random.uniform(shape=(b, h, 20, d))
@@ -98,7 +98,7 @@ class TestModels(unittest.TestCase):
         # alternating prompt/prefill with generation
         d = 4
         h = 2
-        cache = RotatingKVCache(max_size=18, step=4)
+        cache = RotatingKVCache(max_size=18)
 
         x = mx.random.uniform(shape=(1, h, 8, d))
         k, v = cache.update_and_fetch(x, x)
@@ -175,6 +175,49 @@ class TestModels(unittest.TestCase):
         sums = mask.sum(axis=1)
         self.assertTrue(mx.array_equal(sums, expected_sums))
 
+    def test_llama_model_sliding_attention(self):
+        from mlx_lm.models import llama
+
+        args = llama.ModelArgs(
+            model_type="llama",
+            hidden_size=64,
+            num_hidden_layers=4,
+            intermediate_size=256,
+            num_attention_heads=8,
+            num_key_value_heads=4,
+            rms_norm_eps=1e-5,
+            vocab_size=128,
+            sliding_window=4,
+            layer_types=[
+                "full_attention",
+                "sliding_attention",
+                "sliding_attention",
+                "full_attention",
+            ],
+            tie_word_embeddings=False,
+            rope_theta=10000.0,
+        )
+        model = llama.Model(args)
+
+        tokens = mx.array([[1, 2, 3, 4, 5]], dtype=mx.int32)
+        out = model(tokens)
+        mx.eval(out)
+        self.assertEqual(out.shape, (1, 5, args.vocab_size))
+
+        caches = model.make_cache()
+        self.assertIsInstance(caches[0], KVCache)
+        self.assertIsInstance(caches[1], RotatingKVCache)
+        self.assertIsInstance(caches[2], RotatingKVCache)
+        self.assertIsInstance(caches[3], KVCache)
+
+        caches = model.make_cache()
+        step = model(tokens[:, :2], cache=caches)
+        mx.eval(step)
+        step = model(tokens[:, 2:3], cache=caches)
+        mx.eval(step)
+        self.assertEqual(caches[0].offset, 3)
+        self.assertEqual(caches[1].offset, 3)
+
     def test_rope(self):
         rope = rope_utils.initialize_rope(32, base=100, traditional=False)
         self.assertTrue(isinstance(rope, nn.RoPE))
@@ -304,6 +347,35 @@ class TestModels(unittest.TestCase):
             model, args.model_type, args.vocab_size, args.num_hidden_layers
         )
 
+    def test_lfm2_moe(self):
+        from mlx_lm.models import lfm2_moe
+
+        args = lfm2_moe.ModelArgs(
+            model_type="lfm2_moe",
+            hidden_size=1024,
+            intermediate_size=7168,
+            num_hidden_layers=4,
+            num_attention_heads=4,
+            num_key_value_heads=2,
+            norm_eps=1e-5,
+            vocab_size=10_000,
+            full_attn_idxs=[0, 1, 2],
+            rope_theta=10000,
+            max_position_embeddings=1000,
+            conv_bias=True,
+            conv_L_cache=3,
+            moe_intermediate_size=1792,
+            num_dense_layers=2,
+            num_experts=4,
+            num_experts_per_tok=2,
+            norm_topk_prob=True,
+            use_expert_bias=True,
+        )
+        model = lfm2_moe.Model(args)
+        self.model_test_runner(
+            model, args.model_type, args.vocab_size, args.num_hidden_layers
+        )
+
     def test_bitnet(self):
         from mlx_lm.models import bitnet
 
@@ -666,6 +738,19 @@ class TestModels(unittest.TestCase):
             model, args.model_type, args.vocab_size, args.num_hidden_layers
         )
 
+    def test_falcon_h1(self):
+        from mlx_lm.models import falcon_h1
+
+        args = falcon_h1.ModelArgs(
+            model_type="falcon_h1",
+            num_hidden_layers=12,
+            vocab_size=10000,
+        )
+        model = falcon_h1.Model(args)
+        self.model_test_runner(
+            model, args.model_type, args.vocab_size, args.num_hidden_layers
+        )
+
     def test_gpt2(self):
         from mlx_lm.models import gpt2
 
@@ -1650,6 +1735,50 @@ class TestModels(unittest.TestCase):
                 "num_hidden_layers": 4,
                 "vocab_size": 1000,
             },
+            {
+                "model_type": "qwen3_vl_moe",
+                "text_config": {
+                    "model_type": "qwen3_moe",
+                    "hidden_size": 128,
+                    "num_hidden_layers": 4,
+                    "intermediate_size": 256,
+                    "num_attention_heads": 4,
+                    "num_key_value_heads": 2,
+                    "rms_norm_eps": 1e-5,
+                    "head_dim": 32,
+                    "vocab_size": 1000,
+                    "decoder_sparse_step": 1,
+                    "mlp_only_layers": [],
+                    "num_experts_per_tok": 2,
+                    "num_experts": 4,
+                    "moe_intermediate_size": 128,
+                    "rope_theta": 1000,
+                    "max_position_embeddings": 1000,
+                    "tie_word_embeddings": False,
+                    "norm_topk_prob": True,
+                },
+                "num_hidden_layers": 4,
+                "vocab_size": 1000,
+            },
+            {
+                "model_type": "qwen3_vl",
+                "text_config": {
+                    "model_type": "qwen3",
+                    "hidden_size": 128,
+                    "num_hidden_layers": 4,
+                    "intermediate_size": 256,
+                    "num_attention_heads": 4,
+                    "num_key_value_heads": 2,
+                    "rms_norm_eps": 1e-5,
+                    "vocab_size": 1000,
+                    "head_dim": 32,
+                    "max_position_embeddings": 1000,
+                    "tie_word_embeddings": False,
+                    "rope_theta": 1000,
+                },
+                "num_hidden_layers": 4,
+                "vocab_size": 1000,
+            },
             {
                 "model_type": "seed_oss",
                 "hidden_size": 128,
@@ -1768,6 +1897,154 @@ class TestModels(unittest.TestCase):
                 "time_step_limit": (0.01, 10),
                 "time_step_rank": "auto",
             },
+            {
+                "model_type": "olmo3",
+                "num_heads": 8,
+                "head_dim": 16,
+                "vocab_size": 1000,
+                "hidden_size": 128,
+                "intermediate_size": 128,
+                "num_attention_heads": 8,
+                "rope_theta": 1000,
+                "num_hidden_layers": 8,
+                "rms_norm_eps": 1e-4,
+                "sliding_window": 128,
+                "tie_word_embeddings": True,
+                "max_position_embeddings": 1000,
+            },
+            {
+                "model_type": "jamba",
+                "hidden_size": 128,
+                "intermediate_size": 128,
+                "num_hidden_layers": 8,
+                "num_attention_heads": 4,
+                "num_key_value_heads": 2,
+                "attn_layer_offset": 1,
+                "attn_layer_period": 2,
+                "expert_layer_offset": 1,
+                "expert_layer_period": 2,
+                "mamba_d_conv": 4,
+                "mamba_d_state": 128,
+                "mamba_expand": 128,
+                "num_experts": 4,
+                "num_experts_per_tok": 2,
+                "rms_norm_eps": 1e-5,
+                "max_position_embeddings": 1000,
+                "vocab_size": 1000,
+            },
+            {
+                "model_type": "nanochat",
+                "hidden_size": 1280,
+                "num_hidden_layers": 20,
+                "vocab_size": 32,
+                "intermediate_size": 128,
+            },
+            {
+                "model_type": "minimax",
+                "hidden_size": 128,
+                "intermediate_size": 128,
+                "num_attention_heads": 8,
+                "num_key_value_heads": 8,
+                "max_position_embeddings": 1000,
+                "num_experts_per_tok": 2,
+                "num_local_experts": 8,
+                "shared_intermediate_size": 128,
+                "num_hidden_layers": 4,
+                "rms_norm_eps": 1e-4,
+                "rope_theta": 1000,
+                "rotary_dim": 16,
+                "vocab_size": 1000,
+            },
+            {
+                "model_type": "bailing_moe_linear",
+                "hidden_size": 1024,
+                "num_hidden_layers": 4,
+                "intermediate_size": 2048,
+                "moe_intermediate_size": 1024,
+                "num_experts_per_tok": 2,
+                "num_experts": 4,
+                "norm_topk_prob": True,
+                "num_shared_experts": 2,
+                "num_attention_heads": 4,
+                "num_key_value_heads": 4,
+                "rms_norm_eps": 1e-5,
+                "vocab_size": 10_000,
+                "rope_theta": 1000,
+                "first_k_dense_replace": 0,
+                "layer_group_size": 2,
+                "group_norm_size": 1,
+                "max_position_embeddings": 1000,
+            },
+            {
+                "model_type": "kimi_linear",
+                "vocab_size": 1000,
+                "hidden_size": 128,
+                "num_hidden_layers": 4,
+                "num_attention_heads": 8,
+                "num_key_value_heads": 4,
+                "intermediate_size": 128,
+                "head_dim": 32,
+                "rope_theta": 100.0,
+                "rms_norm_eps": 1e-6,
+                "linear_attn_config": {
+                    "num_heads": 8,
+                    "head_dim": 32,
+                    "kda_layers": [1],
+                },
+                "model_max_length": 1000,
+                "num_experts": 2,
+                "moe_intermediate_size": 128,
+                "kv_lora_rank": 8,
+            },
+            {
+                "model_type": "afmoe",
+                "vocab_size": 1000,
+                "hidden_size": 128,
+                "num_hidden_layers": 4,
+                "num_attention_heads": 8,
+                "num_key_value_heads": 4,
+                "intermediate_size": 128,
+                "head_dim": 32,
+                "rope_theta": 100.0,
+                "layer_types": [
+                    "full_attention",
+                    "sliding_attention",
+                    "sliding_attention",
+                    "full_attention",
+                ],
+                "num_experts": 4,
+                "num_experts_per_tok": 2,
+                "moe_intermediate_size": 128,
+            },
+            {
+                "model_type": "deepseek_v32",
+                "vocab_size": 1024,
+                "hidden_size": 128,
+                "intermediate_size": 256,
+                "moe_intermediate_size": 256,
+                "num_hidden_layers": 4,
+                "num_attention_heads": 4,
+                "num_key_value_heads": 2,
+                "n_routed_experts": 4,
+                "n_group": 2,
+                "topk_group": 1,
+                "num_experts_per_tok": 2,
+                "n_shared_experts": 1,
+                "kv_lora_rank": 4,
+                "q_lora_rank": 4,
+                "qk_rope_head_dim": 32,
+                "v_head_dim": 16,
+                "qk_nope_head_dim": 32,
+                "rope_scaling": {
+                    "beta_fast": 32,
+                    "beta_slow": 1,
+                    "factor": 40,
+                    "mscale": 1.0,
+                    "mscale_all_dim": 1.0,
+                    "original_max_position_embeddings": 4096,
+                    "type": "yarn",
+                },
+            },
         ]
         for config in test_configs:
             model_type = config["model_type"]
@@ -1849,9 +2126,9 @@ class TestModels(unittest.TestCase):
         self.assertTrue(mx.allclose(out_state, out_state_m, atol=1e-4, rtol=1e-4))
 
     def test_gated_delta(self):
+        mx.random.seed(0)
         for B in [1, 2]:
             for T in [1, 2]:
-                B = 1
                 Hk = 16
                 Hv = 32
                 Dk = 128
@@ -1860,14 +2137,45 @@ class TestModels(unittest.TestCase):
                 q = mx.random.normal(shape=(B, T, Hk, Dk))
                 k = mx.random.normal(shape=(B, T, Hk, Dk))
                 v = mx.random.normal(shape=(B, T, Hv, Dv))
-                g = mx.random.normal(shape=(B, T, Hv))
-                beta = mx.random.normal(shape=(B, T, Hv))
+                g = mx.random.uniform(shape=(B, T, Hv))
+                beta = mx.random.uniform(shape=(B, T, Hv))
                 state = mx.random.normal(shape=(B, Hv, Dk, Dv))
 
                 y_op, st_op = gated_delta_ops(q, k, v, g, beta, state)
                 y_c, st_c = gated_delta_kernel(q, k, v, g, beta, state)
                 self.assertTrue(mx.allclose(y_op, y_c, rtol=1e-4, atol=1e-4))
-                self.assertTrue(mx.allclose(st_op, st_c, rtol=1e-4, atol=1e-3))
+                self.assertTrue(mx.allclose(st_op, st_c, rtol=1e-4, atol=1e-4))
+
+    def test_gated_delta_masked(self):
+        B = 1
+        T = 3
+        Hk = 16
+        Hv = 32
+        Dk = 128
+        Dv = 128
+
+        mx.random.seed(0)
+        q = mx.random.normal(shape=(B, T, Hk, Dk))
+        k = mx.random.normal(shape=(B, T, Hk, Dk))
+        v = mx.random.normal(shape=(B, T, Hv, Dv))
+        g = mx.random.normal(shape=(B, T, Hv))
+        mask = mx.array([[False, True, True]])
+        beta = mx.random.normal(shape=(B, T, Hv))
+        state = mx.random.normal(shape=(B, Hv, Dk, Dv))
+
+        y_gt, st_gt = gated_delta_ops(
+            q[:, 1:],
+            k[:, 1:],
+            v[:, 1:],
+            g[:, 1:],
+            beta[:, 1:],
+            state,
+        )
+        for fn in [gated_delta_ops, gated_delta_kernel]:
+            y, st = fn(q, k, v, g, beta, state, mask)
+            y = y[:, 1:]
+            self.assertTrue(mx.allclose(y, y_gt, rtol=1e-4, atol=1e-4))
+            self.assertTrue(mx.allclose(st, st_gt, rtol=1e-4, atol=1e-3))
 
 
 if __name__ == "__main__":

tests/test_prompt_cache.py

@@ -11,6 +11,7 @@ from mlx_lm.generate import generate_step
 from mlx_lm.models.base import create_attention_mask, create_causal_mask
 from mlx_lm.models.cache import (
     BatchKVCache,
+    BatchRotatingKVCache,
     CacheList,
     ChunkedKVCache,
     KVCache,
@@ -391,7 +392,7 @@ class TestPromptCache(unittest.TestCase):
         kv = mx.zeros((1, 1, 10, 32))
         cache.update_and_fetch(kv, kv)
         mask = cache.make_mask(3, window_size=5)
-        self.assertEqual(mask.shape, (3, 11))
+        self.assertEqual(mask.shape, (3, 10))
         self.assertTrue(mx.all(mask.sum(axis=-1) == 5))
         for i in range(3):
             s = 11 - 3 + i
@@ -405,7 +406,7 @@ class TestPromptCache(unittest.TestCase):
         self.assertEqual(mask, None)
 
         mask = cache.make_mask(1, window_size=5)
-        self.assertEqual(mask.squeeze(1).tolist(), [True] + [False] * 3 + [True] * 4)
+        self.assertEqual(mask.tolist(), [True] + [False] * 3 + [True] * 4)
         cmask = create_attention_mask(mx.zeros((1, 1, 32)), cache, window_size=5)
         self.assertTrue(mx.array_equal(cmask, mask))
 
@@ -413,9 +414,7 @@ class TestPromptCache(unittest.TestCase):
         cache.update_and_fetch(kv, kv)
 
         mask = cache.make_mask(1, window_size=5)
-        self.assertEqual(
-            mask.squeeze(1).tolist(), [True] * 2 + [False] * 3 + [True] * 3
-        )
+        self.assertEqual(mask.tolist(), [True] * 2 + [False] * 3 + [True] * 3)
         cmask = create_attention_mask(mx.zeros((1, 1, 32)), cache, window_size=5)
         self.assertTrue(mx.array_equal(cmask, mask))
 
@@ -460,6 +459,115 @@ class TestPromptCache(unittest.TestCase):
         self.assertEqual(cache_a.offset.tolist(), [6, 7, 6, 1, 4])
         self.assertEqual(cache_a.left_padding.tolist(), [2, 1, 2, 7, 4])
 
+    def test_batch_rotating_kv_cache(self):
+        cache = BatchRotatingKVCache(max_size=4, left_padding=[2, 0])
+        mask = cache.make_mask(4)
+        self.assertFalse(mx.any(mask[0, 0, 0, :]))
+        self.assertTrue(
+            mx.array_equal(mask[1, 0, 0, :], mx.array([True, False, False, False]))
+        )
+
+        # Batch update works
+        k, v = mx.zeros((2, 1, 4, 8)), mx.zeros((2, 1, 4, 8))
+        k, v = cache.update_and_fetch(k, v)
+
+        mask = cache.make_mask(4)
+        k, v = mx.zeros((2, 1, 4, 8)), mx.zeros((2, 1, 4, 8))
+        k, v = cache.update_and_fetch(k, v)
+        self.assertEqual(mask.shape[-2:], (4, k.shape[2]))
+        self.assertEqual(
+            mask[0, 0, 0, :].tolist(), [False, True, True, True, False, False, False]
+        )
+
+        # Single query update works
+        cache = BatchRotatingKVCache(max_size=4, left_padding=[2, 0])
+        k, v = mx.zeros((2, 1, 4, 8)), mx.zeros((2, 1, 4, 8))
+        k, v = cache.update_and_fetch(k, v)
+
+        mask = cache.make_mask(1)
+        k, v = mx.zeros((2, 1, 1, 8)), mx.zeros((2, 1, 1, 8))
+
+        k, v = cache.update_and_fetch(k, v)
+        self.assertEqual(mask.shape[-2:], (1, k.shape[2]))
+        self.assertEqual(mask[0, 0, 0].tolist(), [True, False, True, True])
+        self.assertEqual(mask[1, 0, 0].tolist(), [True, True, True, True])
+
+        # Check filtering
+        cache = BatchRotatingKVCache(max_size=4, left_padding=[2, 0, 3])
+        k, v = mx.zeros((3, 1, 3, 8)), mx.zeros((3, 1, 3, 8))
+        cache.update_and_fetch(k, v)
+        cache.filter(mx.array([1]))
+        self.assertEqual(cache.keys.shape, (1, 1, 3, 8))
+
+        # Check extend
+        cache = BatchRotatingKVCache(max_size=4, left_padding=[2, 1])
+        other = BatchRotatingKVCache(max_size=4, left_padding=[2, 2])
+        k, v = mx.zeros((2, 1, 5, 8)), mx.zeros((2, 1, 5, 8))
+        cache.update_and_fetch(k, v)
+        other.update_and_fetch(k, v)
+        k, v = mx.zeros((2, 1, 1, 8)), mx.zeros((2, 1, 1, 8))
+        cache.update_and_fetch(k, v)
+        cache.extend(other)
+
+        # Check mask when going from prompt -> extend -> prompt
+        cache = BatchRotatingKVCache(max_size=8, left_padding=[4])
+        k, v = mx.zeros((1, 1, 8, 8)), mx.zeros((1, 1, 8, 8))
+        cache.update_and_fetch(k, v)
+
+        mask = cache.make_mask(1)
+        self.assertEqual(
+            mask.squeeze().tolist(), [True, False, False, False, True, True, True, True]
+        )
+
+        k, v = mx.zeros((1, 1, 1, 8)), mx.zeros((1, 1, 1, 8))
+        cache.update_and_fetch(k, v)
+
+        mask = cache.make_mask(2)
+        expected = mx.array(
+            [
+                [False, False, False, True, True, True, True, True, False],
+                [False, False, False, True, True, True, True, True, True],
+            ]
+        )
+        self.assertTrue(mx.array_equal(mask.squeeze(), expected))
+
+    def test_save_load_batch_caches(self):
+        cache_file = os.path.join(self.test_dir, "prompt_cache.safetensors")
+
+        cache = [
+            MambaCache(left_padding=[1, 2]),
+            BatchKVCache(left_padding=[1, 2]),
+            BatchRotatingKVCache(max_size=10, left_padding=[1, 2]),
+        ]
+        for c in cache:
+            if isinstance(c, MambaCache):
+                c[0] = mx.random.uniform(shape=(4, 4, 4))
+                c[1] = mx.random.uniform(shape=(4, 4, 4))
+            else:
+                x = mx.random.uniform(shape=(4, 4, 7, 4))
+                y = mx.random.uniform(shape=(4, 4, 7, 4))
+                c.update_and_fetch(x, y)
+
+        save_prompt_cache(cache_file, cache)
+        loaded_cache = load_prompt_cache(cache_file)
+        left_padding = mx.array([1, 2])
+        for c, lc in zip(cache, loaded_cache):
+            self.assertTrue(mx.array_equal(c.left_padding, left_padding))
+
+    def test_rotating_cache_updates(self):
+        cache = RotatingKVCache(max_size=8)
+        k = v = mx.zeros((1, 1, 10, 1))
+        cache.update_and_fetch(k, v)
+
+        for _ in range(3):
+            k = v = mx.zeros((1, 1, 1, 1))
+            cache.update_and_fetch(k, v)
+
+        k = v = mx.zeros((1, 1, 3, 1))
+        k, v = cache.update_and_fetch(k, v)
+        self.assertEqual(k.shape[2], 10)
+        self.assertEqual(v.shape[2], 10)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_server.py

@@ -6,9 +6,11 @@ import json
 import threading
 import unittest
 
+import mlx.core as mx
 import requests
 
-from mlx_lm.server import APIHandler
+from mlx_lm.models.cache import KVCache
+from mlx_lm.server import APIHandler, LRUPromptCache, ResponseGenerator
 from mlx_lm.utils import load
 
 
@@ -29,6 +31,7 @@ class DummyModelProvider:
                 "chat_template": None,
                 "use_default_chat_template": False,
                 "trust_remote_code": False,
+                "draft_model": None,
                 "num_draft_tokens": 3,
                 "temp": 0.0,
                 "top_p": 1.0,
@@ -43,6 +46,7 @@ class DummyModelProvider:
             # Use the same model as the draft model for testing
             self.draft_model, _ = load(HF_MODEL_PATH)
             self.draft_model_key = HF_MODEL_PATH
+            self.cli_args.draft_model = HF_MODEL_PATH
 
     def load(self, model, adapter=None, draft_model=None):
         assert model in ["default_model", "chat_model"]
@@ -52,11 +56,13 @@ class DummyModelProvider:
 class TestServer(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
-        cls.model_provider = DummyModelProvider()
+        cls.response_generator = ResponseGenerator(
+            DummyModelProvider(), LRUPromptCache()
+        )
         cls.server_address = ("localhost", 0)
         cls.httpd = http.server.HTTPServer(
             cls.server_address,
-            lambda *args, **kwargs: APIHandler(cls.model_provider, *args, **kwargs),
+            lambda *args, **kwargs: APIHandler(cls.response_generator, *args, **kwargs),
         )
         cls.port = cls.httpd.server_port
         cls.server_thread = threading.Thread(target=cls.httpd.serve_forever)
@@ -68,6 +74,7 @@ class TestServer(unittest.TestCase):
         cls.httpd.shutdown()
         cls.httpd.server_close()
         cls.server_thread.join()
+        cls.response_generator.stop_and_join()
 
     def test_handle_completions(self):
         url = f"http://localhost:{self.port}/v1/completions"
@@ -198,11 +205,13 @@ class TestServer(unittest.TestCase):
 class TestServerWithDraftModel(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
-        cls.model_provider = DummyModelProvider(with_draft=True)
+        cls.response_generator = ResponseGenerator(
+            DummyModelProvider(with_draft=True), LRUPromptCache()
+        )
         cls.server_address = ("localhost", 0)
         cls.httpd = http.server.HTTPServer(
             cls.server_address,
-            lambda *args, **kwargs: APIHandler(cls.model_provider, *args, **kwargs),
+            lambda *args, **kwargs: APIHandler(cls.response_generator, *args, **kwargs),
         )
         cls.port = cls.httpd.server_port
         cls.server_thread = threading.Thread(target=cls.httpd.serve_forever)
@@ -214,6 +223,7 @@ class TestServerWithDraftModel(unittest.TestCase):
         cls.httpd.shutdown()
         cls.httpd.server_close()
         cls.server_thread.join()
+        cls.response_generator.stop_and_join()
 
     def test_handle_completions_with_draft_model(self):
         url = f"http://localhost:{self.port}/v1/completions"
@@ -339,182 +349,6 @@ class TestServerWithDraftModel(unittest.TestCase):
         self.assertIsNotNone(second_response_body["choices"][0]["message"]["content"])
 
 
-# --- Tests for get_prompt_cache ---
-
-from unittest.mock import MagicMock, patch
-
-from mlx_lm.server import PromptCache
-
-
-class TestGetPromptCache(unittest.TestCase):
-
-    def setUp(self):
-        """Set up mocks and a handler instance for each test."""
-        self.mock_model_provider = MagicMock()
-        # Simulate tokenizer needed for decoding in original debug logs (though not strictly needed for cache logic)
-        self.mock_model_provider.tokenizer = MagicMock()
-        self.mock_model_provider.tokenizer.decode = lambda x: f"decoded({x})"
-        self.mock_model_provider.model_key = ("model_v1", None, None)
-        self.mock_model_provider.draft_model = None  # Start without draft model
-
-        # --- Prevent BaseHTTPRequestHandler.__init__ from running ---
-        # It tries to handle a request immediately, which fails with mocks.
-        # We only need the APIHandler instance with its attributes set.
-        with patch(
-            "http.server.BaseHTTPRequestHandler.__init__", lambda *args, **kwargs: None
-        ):
-            # APIHandler init still requires args for BaseHTTPRequestHandler signature,
-            # but they won't be used by the patched __init__.
-            mock_request = MagicMock()
-            mock_client_address = ("127.0.0.1", 8080)
-            mock_server = MagicMock()
-
-            self.prompt_cache_instance = PromptCache()
-            self.handler = APIHandler(
-                self.mock_model_provider,
-                mock_request,
-                mock_client_address,
-                mock_server,
-                prompt_cache=self.prompt_cache_instance,  # Inject our cache instance
-            )
-            # Manually set attributes usually set by APIHandler.__init__ if needed
-            # self.handler.created = MagicMock()
-            # self.handler.system_fingerprint = MagicMock()
-            # (Not strictly necessary for get_prompt_cache testing)
-
-    @patch("mlx_lm.server.make_prompt_cache")
-    def test_initial_request_empty_cache(self, mock_make_cache):
-        """Test first request when the cache is empty."""
-        mock_make_cache.return_value = "new_cache_obj"
-        prompt = [1, 2, 3]
-
-        processed_prompt = self.handler.get_prompt_cache(prompt)
-
-        self.assertEqual(processed_prompt, [1, 2, 3])
-        self.assertEqual(self.handler.prompt_cache.tokens, [1, 2, 3])
-        self.assertEqual(self.handler.prompt_cache.cache, "new_cache_obj")
-        self.assertEqual(self.handler.prompt_cache.model_key, ("model_v1", None, None))
-        mock_make_cache.assert_called_once()
-
-    @patch("mlx_lm.server.trim_prompt_cache")
-    @patch("mlx_lm.server.can_trim_prompt_cache", return_value=True)
-    def test_identical_request_full_hit(self, mock_can_trim, mock_trim_cache):
-        """Test when the new prompt is identical to the cached one."""
-        self.handler.prompt_cache.tokens = [1, 2, 3]
-        self.handler.prompt_cache.model_key = ("model_v1", None, None)
-        self.handler.prompt_cache.cache = "existing_cache_obj"
-        prompt = [1, 2, 3]
-
-        # Mock common_prefix_len to return the full length
-        with patch("mlx_lm.server.common_prefix_len", return_value=3):
-            processed_prompt = self.handler.get_prompt_cache(prompt)
-
-        mock_trim_cache.assert_called_once_with("existing_cache_obj", 1)
-        self.assertEqual(processed_prompt, [3])
-        self.assertEqual(self.handler.prompt_cache.tokens, [1, 2, 3])
-
-    def test_cache_is_prefix(self):
-        """Test when the cached prompt is a prefix of the new prompt."""
-        self.handler.prompt_cache.tokens = [1, 2, 3]
-        self.handler.prompt_cache.model_key = ("model_v1", None, None)
-        self.handler.prompt_cache.cache = "existing_cache_obj"
-        prompt = [1, 2, 3, 4, 5]
-
-        with patch("mlx_lm.server.common_prefix_len", return_value=3):
-            processed_prompt = self.handler.get_prompt_cache(prompt)
-
-        # Should process the suffix, cache tokens updated
-        self.assertEqual(processed_prompt, [4, 5])
-        self.assertEqual(self.handler.prompt_cache.tokens, [1, 2, 3, 4, 5])
-        self.assertEqual(self.handler.prompt_cache.cache, "existing_cache_obj")
-
-    @patch("mlx_lm.server.trim_prompt_cache")
-    @patch("mlx_lm.server.can_trim_prompt_cache", return_value=True)
-    def test_partial_match_trim_success(self, mock_can_trim, mock_trim_cache):
-        """Test partial match where cache is longer and trimming succeeds."""
-        self.handler.prompt_cache.tokens = [1, 2, 3, 4, 5]
-        self.handler.prompt_cache.model_key = ("model_v1", None, None)
-        self.handler.prompt_cache.cache = "existing_cache_obj"
-        prompt = [1, 2, 3, 6, 7]  # Diverges after token 3
-
-        with patch("mlx_lm.server.common_prefix_len", return_value=3):
-            processed_prompt = self.handler.get_prompt_cache(prompt)
-
-        # Should process the new suffix, cache trimmed and updated
-        self.assertEqual(processed_prompt, [6, 7])
-        self.assertEqual(self.handler.prompt_cache.tokens, [1, 2, 3, 6, 7])
-        mock_can_trim.assert_called_once_with("existing_cache_obj")
-        # Called with cache object and num_to_trim (5 - 3 = 2)
-        mock_trim_cache.assert_called_once_with("existing_cache_obj", 2)
-        self.assertEqual(
-            self.handler.prompt_cache.cache, "existing_cache_obj"
-        )  # Cache obj itself isn't changed by mock
-
-    @patch("mlx_lm.server.make_prompt_cache")
-    @patch("mlx_lm.server.trim_prompt_cache")
-    @patch("mlx_lm.server.can_trim_prompt_cache", return_value=False)
-    def test_partial_match_trim_fail(
-        self, mock_can_trim, mock_trim_cache, mock_make_cache
-    ):
-        """Test partial match where cache is longer but trimming fails."""
-        mock_make_cache.return_value = "new_cache_obj_on_reset"
-        self.handler.prompt_cache.tokens = [1, 2, 3, 4, 5]
-        self.handler.prompt_cache.model_key = ("model_v1", None, None)
-        self.handler.prompt_cache.cache = "existing_cache_obj"
-        prompt = [1, 2, 3, 6, 7]  # Diverges after token 3
-
-        with patch("mlx_lm.server.common_prefix_len", return_value=3):
-            processed_prompt = self.handler.get_prompt_cache(prompt)
-
-        # Should process the full prompt, cache reset
-        self.assertEqual(processed_prompt, [1, 2, 3, 6, 7])
-        self.assertEqual(self.handler.prompt_cache.tokens, [1, 2, 3, 6, 7])
-        mock_can_trim.assert_called_once_with("existing_cache_obj")
-        mock_trim_cache.assert_not_called()
-        mock_make_cache.assert_called_once()  # Cache was reset
-        self.assertEqual(self.handler.prompt_cache.cache, "new_cache_obj_on_reset")
-
-    @patch("mlx_lm.server.make_prompt_cache")
-    def test_no_common_prefix(self, mock_make_cache):
-        """Test when there is no common prefix between cache and prompt."""
-        mock_make_cache.return_value = "new_cache_obj"
-        self.handler.prompt_cache.tokens = [1, 2, 3]
-        self.handler.prompt_cache.model_key = ("model_v1", None, None)
-        self.handler.prompt_cache.cache = "existing_cache_obj"
-        prompt = [4, 5, 6]
-
-        with patch("mlx_lm.server.common_prefix_len", return_value=0):
-            processed_prompt = self.handler.get_prompt_cache(prompt)
-
-        # Should process the full prompt, cache reset
-        self.assertEqual(processed_prompt, [4, 5, 6])
-        self.assertEqual(self.handler.prompt_cache.tokens, [4, 5, 6])
-        mock_make_cache.assert_called_once()
-        self.assertEqual(self.handler.prompt_cache.cache, "new_cache_obj")
-
-    @patch("mlx_lm.server.make_prompt_cache")
-    def test_model_changed(self, mock_make_cache):
-        """Test cache reset when the model key changes."""
-        mock_make_cache.return_value = "new_cache_obj_model_change"
-        self.handler.prompt_cache.tokens = [1, 2, 3]
-        self.handler.prompt_cache.model_key = ("model_v1", None, None)  # Original key
-        self.handler.prompt_cache.cache = "existing_cache_obj"
-
-        # Simulate model provider having a new key
-        self.mock_model_provider.model_key = ("model_v2", None, None)
-        prompt = [1, 2, 3, 4]
-
-        # No need to mock common_prefix_len, model check happens first
-        processed_prompt = self.handler.get_prompt_cache(prompt)
-
-        # Should process the full prompt, cache reset
-        self.assertEqual(processed_prompt, [1, 2, 3, 4])
-        self.assertEqual(self.handler.prompt_cache.tokens, [1, 2, 3, 4])
-        mock_make_cache.assert_called_once()
-        self.assertEqual(self.handler.prompt_cache.cache, "new_cache_obj_model_change")
-        self.assertEqual(self.handler.prompt_cache.model_key, ("model_v2", None, None))
-
-
 class TestKeepalive(unittest.TestCase):
 
     def test_keepalive_callback(self):
@@ -565,5 +399,78 @@ class TestKeepalive(unittest.TestCase):
             self.fail(f"Callback should handle BrokenPipeError: {e}")
 
 
+class TestLRUPromptCache(unittest.TestCase):
+
+    def test_caching(self):
+        cache = LRUPromptCache(max_size=10)
+
+        def get_kv(n):
+            keys = mx.arange(n).reshape(1, 1, n, 1)
+            return keys, keys
+
+        model = ("test", None, None)
+        tokens = [10] * 24
+
+        c, t = cache.fetch_nearest_cache(model, tokens)
+        self.assertTrue(c is None)
+        self.assertEqual(t, tokens)
+
+        c = [KVCache()]
+        c[0].update_and_fetch(*get_kv(24))
+        cache.insert_cache(model, t, c)
+
+        tokens = tokens + [20] * 5
+        c, t = cache.fetch_nearest_cache(model, tokens)
+        k, v = c[0].state
+        self.assertTrue((k == v).all().item())
+        self.assertTrue((k.flatten() == mx.arange(24)).all().item())
+        self.assertEqual(t, [20] * 5)
+        self.assertEqual(len(cache._lru), 0)
+
+        tokens = tokens + [30] * 3
+        c[0].update_and_fetch(*get_kv(8))
+        cache.insert_cache(model, tokens, c)
+
+        tokens = tokens[:26] + [40] * 8
+        c, t = cache.fetch_nearest_cache(model, tokens)
+        k, v = c[0].state
+        self.assertTrue((k == v).all().item())
+        self.assertTrue(
+            (k.flatten() == mx.concatenate([mx.arange(24), mx.arange(2)])).all().item()
+        )
+        self.assertEqual(t, [40] * 8)
+        self.assertEqual(len(cache._lru), 1)
+
+    def test_lru(self):
+        cache = LRUPromptCache(max_size=2)
+        model = ("test", None, None)
+        cache.insert_cache(model, [1, 2], ["test1"])
+        cache.insert_cache(model, [1, 2], ["test1"])
+
+        c, t = cache.fetch_nearest_cache(model, [1, 2])
+        self.assertEqual(c, ["test1"])
+        self.assertEqual(t, [])
+        c, t = cache.fetch_nearest_cache(model, [1, 2])
+        self.assertEqual(c, ["test1"])
+        self.assertEqual(t, [])
+        c, t = cache.fetch_nearest_cache(model, [1, 2])
+        self.assertEqual(c, None)
+        self.assertEqual(t, [1, 2])
+
+        cache.insert_cache(model, [1, 2], ["test1"])
+        cache.insert_cache(model, [2, 3], ["test2"])
+        cache.insert_cache(model, [3, 4], ["test3"])
+
+        c, t = cache.fetch_nearest_cache(model, [1, 2])
+        self.assertEqual(c, None)
+        self.assertEqual(t, [1, 2])
+        c, t = cache.fetch_nearest_cache(model, [2, 3])
+        self.assertEqual(c, ["test2"])
+        self.assertEqual(t, [])
+        c, t = cache.fetch_nearest_cache(model, [3, 4])
+        self.assertEqual(c, ["test3"])
+        self.assertEqual(t, [])
+
+
 if __name__ == "__main__":
     unittest.main()

tests/test_tokenizers.py

@@ -9,27 +9,12 @@ from mlx_lm.tokenizer_utils import (
     BPEStreamingDetokenizer,
     NaiveStreamingDetokenizer,
     SPMStreamingDetokenizer,
-    load_tokenizer,
 )
+from mlx_lm.utils import load_tokenizer
 
 
 class TestTokenizers(unittest.TestCase):
 
-    def download_tokenizer(self, repo):
-        path = Path(
-            snapshot_download(
-                repo_id=repo,
-                allow_patterns=[
-                    "tokenizer.json",
-                    "tokenizer_config.json",
-                    "special_tokens_map.json",
-                    "tokenizer.model",
-                    "chat_template.jinja",
-                ],
-            )
-        )
-        return load_tokenizer(path)
-
     def check_tokenizer(self, tokenizer):
         def check(tokens):
             expected_text = tokenizer.decode(tokens)
@@ -77,19 +62,19 @@ class TestTokenizers(unittest.TestCase):
         ]
         for tokenizer_repo, expected_detokenizer in tokenizer_repos:
             with self.subTest(tokenizer=tokenizer_repo):
-                tokenizer = self.download_tokenizer(tokenizer_repo)
+                tokenizer = load_tokenizer(tokenizer_repo)
                 tokenizer.decode([0, 1, 2])
                 self.assertTrue(isinstance(tokenizer.detokenizer, expected_detokenizer))
                 self.check_tokenizer(tokenizer)
 
         # Try one with a naive detokenizer
-        tokenizer = self.download_tokenizer("mlx-community/Llama-3.2-1B-Instruct-4bit")
+        tokenizer = load_tokenizer("mlx-community/Llama-3.2-1B-Instruct-4bit")
         tokenizer._detokenizer = NaiveStreamingDetokenizer(tokenizer)
         self.check_tokenizer(tokenizer)
 
     def test_special_tokens(self):
         tokenizer_repo = "mlx-community/DeepSeek-Coder-V2-Lite-Instruct-4bit-mlx"
-        tokenizer = self.download_tokenizer(tokenizer_repo)
+        tokenizer = load_tokenizer(tokenizer_repo)
 
         detokenizer = tokenizer.detokenizer
         detokenizer.reset()
@@ -100,18 +85,18 @@ class TestTokenizers(unittest.TestCase):
 
     def test_tool_calling(self):
         tokenizer_repo = "mlx-community/Qwen3-4B-4bit"
-        tokenizer = self.download_tokenizer(tokenizer_repo)
+        tokenizer = load_tokenizer(tokenizer_repo)
         self.assertTrue(tokenizer.has_tool_calling)
         self.assertEqual(tokenizer.tool_call_start, "<tool_call>")
         self.assertEqual(tokenizer.tool_call_end, "</tool_call>")
 
         tokenizer_repo = "mlx-community/Llama-3.2-1B-Instruct-4bit"
-        tokenizer = self.download_tokenizer(tokenizer_repo)
+        tokenizer = load_tokenizer(tokenizer_repo)
         self.assertFalse(tokenizer.has_tool_calling)
 
     def test_thinking(self):
         tokenizer_repo = "mlx-community/Qwen3-4B-4bit"
-        tokenizer = self.download_tokenizer(tokenizer_repo)
+        tokenizer = load_tokenizer(tokenizer_repo)
         self.assertTrue(tokenizer.has_thinking)
         self.assertEqual(tokenizer.think_start, "<think>")
         self.assertEqual(tokenizer.think_end, "</think>")

tests/test_tuner_trainer.py

@@ -19,47 +19,35 @@ class MockDistributedGroup:
         return self._size
 
 
-class MockDistributed:
-    def __init__(self):
-        self.rank = 0
-        self.size = 1
-
-    def init(self):
-        return MockDistributedGroup(self.rank, self.size)
-
-
 class TestTunerTrainer(unittest.TestCase):
     def test_iterate_batches_ddp(self):
-        olddist = mx.distributed
-        try:
-            mx.distributed = MockDistributed()
+        group = MockDistributedGroup(0, 1)
 
-            def run(rank, size, batch):
-                mx.distributed.rank = rank
-                mx.distributed.size = size
+        def run(rank, size, batch):
+            group._rank = rank
+            group._size = size
 
-                data = mx.arange(128).reshape(-1, 1).tolist()
-                data = [(d, 0) for d in data]
+            data = mx.arange(128).reshape(-1, 1).tolist()
+            data = [(d, 0) for d in data]
 
-                samples = set()
-                for i, (b, l) in enumerate(iterate_batches(data, batch, 1)):
-                    samples.add(tuple(mx.flatten(b).tolist()))
+            samples = set()
+            for i, (b, l) in enumerate(
+                iterate_batches(data, batch, 1, comm_group=group)
+            ):
+                samples.add(tuple(mx.flatten(b).tolist()))
 
-                ref_batches = mx.arange(128).reshape(-1, batch).tolist()
-                for b in ref_batches:
-                    self.assertTrue(tuple(b[rank::size]) in samples)
+            ref_batches = mx.arange(128).reshape(-1, batch).tolist()
+            for b in ref_batches:
+                self.assertTrue(tuple(b[rank::size]) in samples)
 
-            run(0, 1, 4)
-            run(0, 1, 8)
-            run(0, 2, 8)
-            run(1, 2, 8)
-            run(0, 4, 8)
-            run(1, 4, 8)
-            run(2, 4, 8)
-            run(3, 4, 8)
-
-        finally:
-            mx.distributed = olddist
+        run(0, 1, 4)
+        run(0, 1, 8)
+        run(0, 2, 8)
+        run(1, 2, 8)
+        run(0, 4, 8)
+        run(1, 4, 8)
+        run(2, 4, 8)
+        run(3, 4, 8)
 
 
 if __name__ == "__main__":

tests/test_utils.py

@@ -27,7 +27,10 @@ class TestUtils(unittest.TestCase):
         cls.test_dir_fid.cleanup()
 
     def test_load(self):
+        from mlx_lm.models.qwen2 import Model as Qwen2Model
+
         model, _ = utils.load(HF_MODEL_PATH)
+        self.assertIsInstance(model, Qwen2Model)
 
         model_lazy, _ = utils.load(HF_MODEL_PATH, lazy=True)
 
@@ -91,6 +94,35 @@ class TestUtils(unittest.TestCase):
         self.assertEqual(model.layers[0].mlp.up_proj.scales.dtype, mx.bfloat16)
         self.assertEqual(model.layers[-1].mlp.up_proj.scales.dtype, mx.bfloat16)
 
+    def test_load_model_with_custom_get_classes(self):
+        class CustomQwenModel(nn.Module):
+            def __init__(self, args):
+                super().__init__()
+                self.config = args
+                self.custom_attribute = "This is a custom model"
+
+            def load_weights(self, weights, **kwargs):
+                self.qwenWeights = weights
+
+        class CustomQwenConfig:
+            @classmethod
+            def from_dict(cls, config):
+                instance = cls()
+                for k, v in config.items():
+                    setattr(instance, k, v)
+                return instance
+
+        def custom_get_classes(config):
+            return CustomQwenModel, CustomQwenConfig
+
+        model_path = utils.hf_repo_to_path(HF_MODEL_PATH)
+        model, _ = utils.load_model(model_path, get_model_classes=custom_get_classes)
+
+        self.assertIsInstance(model, CustomQwenModel)
+        self.assertTrue(hasattr(model, "custom_attribute"))
+        self.assertEqual(model.custom_attribute, "This is a custom model")
+        self.assertTrue(hasattr(model, "qwenWeights"))
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_utils_load_model.py

@@ -1,51 +0,0 @@
-import unittest
-from pathlib import Path
-
-import mlx.nn as nn
-
-from mlx_lm.models.qwen2 import Model as Qwen2Model
-from mlx_lm.utils import get_model_path, load_model
-
-HF_MODEL_PATH = "mlx-community/Qwen1.5-0.5B-Chat-4bit"
-
-
-class TestLoadModelCustomGetClasses(unittest.TestCase):
-
-    def test_load_model_with_custom_get_classes(self):
-        class CustomQwenModel(nn.Module):
-            def __init__(self, args):
-                super().__init__()
-                self.config = args
-                self.custom_attribute = "This is a custom model"
-
-            def load_weights(self, weights, **kwargs):
-                self.qwenWeights = weights
-
-        class CustomQwenConfig:
-            @classmethod
-            def from_dict(cls, config):
-                instance = cls()
-                for k, v in config.items():
-                    setattr(instance, k, v)
-                return instance
-
-        def custom_get_classes(config):
-            return CustomQwenModel, CustomQwenConfig
-
-        model_path, _ = get_model_path(HF_MODEL_PATH)
-        model, _ = load_model(model_path, get_model_classes=custom_get_classes)
-
-        self.assertIsInstance(model, CustomQwenModel)
-        self.assertTrue(hasattr(model, "custom_attribute"))
-        self.assertEqual(model.custom_attribute, "This is a custom model")
-        self.assertTrue(hasattr(model, "qwenWeights"))
-
-    def test_load_model_with_default_get_classes(self):
-        model_path, _ = get_model_path(HF_MODEL_PATH)
-        model, _ = load_model(model_path)
-
-        self.assertIsInstance(model, Qwen2Model)
-
-
-if __name__ == "__main__":
-    unittest.main()