feat: Refactor granitemoehybrid to support dense and non-hybrid variants (#518)

* feat: Refactor granitemoehybrid to support dense and non-hybrid variants

Written with Claude Code. Initial prompt:

I need to modify the model support implemented in `mlx_lm/models/granitemoehybrid.py` in two ways:

* Support optionally using a dense block in place of MoE. The dense block should look like `mlx_lm/models/granite.py` instead of `mlx_lm/models/granitemoe.py`.

* Support the case where there are no `mamba` layers (ie non-hybrid). This should devolve to exactly `granite.py` or `granitemoe.py` depending on whether the block after attention is dense or MoE.

You can test this using the following two models:

* Dense w/ hybrid: /Users/ghart/models/dmf_models/granite-4.0-h-micro-r250918a
* Dense w/ non-hybrid: /Users/ghart/models/dmf_models/granite-4.0-micro-r250918a

Branch: GraniteFourDense

Signed-off-by: Gabe Goodhart <ghart@us.ibm.com>

* refact: Clean up Claude's code a bit

Branch: GraniteFourDense

Signed-off-by: Gabe Goodhart <ghart@us.ibm.com>

* style: pre-commit format

Branch: GraniteFourDense

Signed-off-by: Gabe Goodhart <ghart@us.ibm.com>

* version bump

---------

Signed-off-by: Gabe Goodhart <ghart@us.ibm.com>
Co-authored-by: Awni Hannun <awni@apple.com>

.circleci/config.yml

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

ACKNOWLEDGMENTS.md

@@ -8,5 +8,17 @@ with a short description of your contribution(s) below. For example:
 MLX LM was developed with contributions from the following individuals:
 
 - Shunta Saito: Added support for PLaMo models.
-- Prince Canuma: Helped add support for `Starcoder2` models.
-- Gökdeniz Gülmez: Added support for the following architectures: OpenBMB's `MiniCPM` and `MiniCPM3`, Kyutai's `Helium`, State-Space's`Mamba v1`, Z.ai & THUKEG's `GLM4`, and Allenai's `OLMoE`; Added support for the following training algorithms: `full-fine-tuning`; Added support for the following other features: `Multiple Optimizers to choose for training`.
+- Gökdeniz Gülmez: Added support for the following architectures: OpenBMB's
+  `MiniCPM` and `MiniCPM3`, Kyutai's `Helium`, State-Space's `Mamba v1`, `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` 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`.

README.md

@@ -52,6 +52,12 @@ options for a command, e.g.:
 mlx_lm.generate -h
 ```
 
+The default model for generation and chat is
+`mlx-community/Llama-3.2-3B-Instruct-4bit`.  You can specify any MLX-compatible
+model with the `--model` flag. Thousands are available in the
+[MLX Community](https://huggingface.co/mlx-community) Hugging Face
+organization.
+
 ### Python API
 
 You can use `mlx-lm` as a module:
@@ -79,7 +85,9 @@ To see a description of all the arguments you can do:
 
 Check out the [generation
 example](https://github.com/ml-explore/mlx-lm/tree/main/mlx_lm/examples/generate_response.py)
-to see how to use the API in more detail.
+to see how to use the API in more detail. Check out the [batch generation
+example](https://github.com/ml-explore/mlx-lm/tree/main/mlx_lm/examples/batch_generate_response.py)
+to see how to efficiently generate continuations for a batch of prompts.
 
 The `mlx-lm` package also comes with functionality to quantize and optionally
 upload models to the Hugging Face Hub.

mlx_lm/LEARNED_QUANTS.md

@@ -1,21 +1,28 @@
 # Learned Quantization 
 
-To reduce the quality loss from quantization MLX LM has two options:
+To reduce the quality loss from quantization MLX LM has several options:
 
 - Distilled Weight Quantization (DWQ)
-- Activation-aware Weight Quantization (AWQ)[^1].
+- Activation-aware Weight Quantization (AWQ)[^1]
+- Dynamic quantization
+- GPT Quantization (GPTQ)[^2]
 
-Both DWQ and AWQ use an example dataset to tune parameters of the model. DWQ
-fine-tunes non-quantized parameters (including quantization scales and biases)
-using the non-quantized model as a teacher. AWQ scales and clips the weights
-prior to quantization. The scaling and clipping values are found with a grid
-search minimizing the distance from the quantized hidden activations to the
-non-quantized hidden activations
+All methods use calibration data to tune parameters or hyper-parameters of the
+model. DWQ fine-tunes non-quantized parameters (including quantization scales
+and biases) using the non-quantized model as a teacher. AWQ scales and clips
+the weights prior to quantization. Dynamic quantization estimates the
+sensitivity of a model's outputs to each layer and uses a higher precision for
+layers which have higher sensitivity. GPTQ finds quantized weights which
+minimize the squared error of each layer's output given the provided input.
+
+Dynamic quantization is the fastest to run. DWQ takes longer but typically
+yields better results. You can also cascade methods. For example a dynamically
+quantized model can be further refined with DWQ.
 
 To get started, first install the requirements:
 
 ```
-pip install mlx-lm[lwq]
+pip install "mlx-lm[train]"
 ```
 
 ### DWQ
@@ -23,7 +30,7 @@ pip install mlx-lm[lwq]
 Use `mlx_lm.dwq` to run DWQ on a given model. For example:
 
 ```bash
-mlx_lm.dwq --model mistralai/Mistral-7B-Instruct-v0.3
+mlx_lm.dwq --model Qwen/Qwen3-0.6B
 ```
 
 Some important options, along with their default values are:
@@ -40,12 +47,63 @@ For a full list of options run:
 mlx_lm.dwq --help
 ```
 
+#### Tips
+
+- DWQ works best distilling to lower precision, anywhere from 2-bit to 4-bit
+  models.
+- Distilling 16-bit precision to 8-bit and even 6-bit often doesn't work well.
+  The loss starts out so low that it's difficult to reduce further.
+- Decreasing the quantization group size (e.g. `--group-size 32`) doubles the
+  number of tunable parameters and can work much better.
+- If the loss is oscillating and not going down consistently, try reducing the
+  learning rate. If it is decreasing but slowly, try increasing the learning
+  rate.
+- As a rule of thumb, lower precision can benefit from a higher learning rate
+  since the loss starts out higher. Conversely, higher precision needs a lower
+  learning rate.
+
+
+#### Memory Use
+
+A few options to reduce memory use for DWQ:
+
+- Distill from an 8-bit model instead of a 16-bit model. The 8-bit
+  models are usually as good as 16-bit precision models.
+- Use a shorter maximum sequence length. The default is 2048. Using
+  `--max-seq-length 512` reduces the memory and still gets good results.
+- Use a smaller batch size, e.g. `--batch-size 1`
+
+### Dynamic Quantization
+
+Use `mlx_lm.dynamic_quant` to generate a dynamic quantization of given model.
+For example:
+
+```bash
+mlx_lm.dynamic_quant --model Qwen/Qwen3-0.6B
+```
+
+The script will estimate the sensitivity for each quantizable layer in the
+model. It will then quantize the model using higher precision (default 5 bits)
+for the more sensitive layers and lower precision (default 4 bits) for the
+rest. The script also saves a JSON file with each layer's sensitivities which
+saves needing to compute it multiple times to make different precision quants
+of the same model.
+
+Some important options are:
+
+- `--target-bpw`: The target bits-per-weight. For a given set of quantization
+  parameters only certain ranges are possible. For example, with the default
+  parameters a BPW in the range `[4.5, 5.5]` is achievable.
+- `--sensitivities`: A path to a precomputed sensitivities file.
+- `--low-bits`: The number of bits to use for the less sensitive layers.
+- `--high-bits`: The number of bits to use for the more sensitive layers.
+
 ### AWQ 
 
 Use `mlx_lm.awq` to run AWQ on a given model. For example:
 
 ```bash
-mlx_lm.awq --model mistralai/Mistral-7B-Instruct-v0.3
+mlx_lm.awq --model Qwen/Qwen3-0.6B
 ```
 
 The script can take anywhere form a few minutes to several hours to run
@@ -66,10 +124,27 @@ For a full list of options run:
 mlx_lm.awq --help
 ```
 
+### GPTQ
+
+Use `mlx_lm.gptq` to run GPTQ on a given model. For example:
+
+```bash
+mlx_lm.awq --model Qwen/Qwen3-0.6B
+```
+
+The script can take anywhere from a few minutes to several hours depending on
+the model size.
+
+Some important options, along with their default values, are:
+
+- `--mlx-path mlx_model`: The location to save the AWQ model.
+- `--bits 4`: Precision of the quantization.
+
+
 ### Evaluate
 
-Once the training script finishes, you can evaluate the quality of the model
-on downstream tasks using `mlx_lm.evaluate`. For example:
+Once the quantization training finishes, you can evaluate the quality of the
+model on downstream tasks using `mlx_lm.evaluate`. For example:
 
 ```bash
 mlx_lm.evaluate \
@@ -90,4 +165,6 @@ mlx_lm.upload \
 
 [^1]: Refer to the [paper](https://arxiv.org/abs/2306.00978)
 and [github repository](https://github.com/mit-han-lab/llm-awq) for more
-details.
+details on AWQ.
+[^2]: Refer to the [paper](https://arxiv.org/abs/2210.17323) for more details
+on GPTQ.

mlx_lm/LORA.md

@@ -26,6 +26,12 @@ LoRA (QLoRA).[^qlora] LoRA fine-tuning works with the following model families:
 
 ## Run
 
+First, make sure you have the training dependenices installed:
+
+```shell
+pip install "mlx-lm[train]"
+```
+
 The main command is `mlx_lm.lora`. To see a full list of command-line options run:
 
 ```shell
@@ -76,6 +82,17 @@ You can specify the output location with `--adapter-path`.
 You can resume fine-tuning with an existing adapter with
 `--resume-adapter-file <path_to_adapters.safetensors>`.
 
+#### Logging
+
+You can log training metrics to Weights & Biases using `--report-to wandb`, or
+to SwanLab using `--report-to swanlab`. Make sure to install the required
+packages beforehand: `pip install wandb` or `pip install swanlab`. You can
+enable both tracking tools simultaneously by separating them with a comma, for
+example: `--report-to wandb,swanlab`.
+
+To specify a project name for the logging tracker, use `--project-name <YOUR
+PROJECT NAME>`.
+
 #### Prompt Masking
 
 The default training computes a loss for every token in the sample. You can
@@ -379,7 +396,7 @@ mlx_lm.lora \
     --train \
     --batch-size 1 \
     --num-layers 4 \
-    --data wikisql
+    --data  mlx-community/wikisql
 ```
 
 The above command on an M1 Max with 32 GB runs at about 250

mlx_lm/MERGE.md

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

mlx_lm/SERVER.md

@@ -54,18 +54,24 @@ curl localhost:8080/v1/chat/completions \
   sequences of tokens on which the generation should stop.
 
 - `max_tokens`: (Optional) An integer specifying the maximum number of tokens
-  to generate. Defaults to `100`.
+  to generate. Defaults to `512`.
 
 - `stream`: (Optional) A boolean indicating if the response should be
   streamed. If true, responses are sent as they are generated. Defaults to
   false.
 
 - `temperature`: (Optional) A float specifying the sampling temperature.
-  Defaults to `1.0`.
+  Defaults to `0.0`.
 
 - `top_p`: (Optional) A float specifying the nucleus sampling parameter.
   Defaults to `1.0`.
 
+- `top_k`: (Optional) An integer specifying the top-k sampling parameter.
+  Defaults to `0` (disabled).
+
+- `min_p`: (Optional) A float specifying the min-p sampling parameter.
+  Defaults to `0.0` (disabled).
+
 - `repetition_penalty`: (Optional) Applies a penalty to repeated tokens.
   Defaults to `1.0`.
 
@@ -92,7 +98,6 @@ curl localhost:8080/v1/chat/completions \
 - `num_draft_tokens`: (Optional) The number of draft tokens the draft model
   should predict at once. Defaults to `3`.
 
-
 ### Response Fields
 
 - `id`: A unique identifier for the chat.

mlx_lm/UPLOAD.md

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

mlx_lm/__init__.py

@@ -7,5 +7,5 @@ from ._version import __version__
 os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1"
 
 from .convert import convert
-from .generate import generate, stream_generate
+from .generate import batch_generate, generate, stream_generate
 from .utils import load

mlx_lm/__main__.py

@@ -5,8 +5,11 @@ import sys
 
 if __name__ == "__main__":
     subcommands = {
-        "awq",
-        "dwq",
+        "quant.awq",
+        "quant.dwq",
+        "quant.dynamic_quant",
+        "quant.gptq",
+        "benchmark",
         "cache_prompt",
         "chat",
         "convert",
@@ -14,7 +17,7 @@ if __name__ == "__main__":
         "fuse",
         "generate",
         "lora",
-        "merge",
+        "perplexity",
         "server",
         "manage",
         "upload",

mlx_lm/_version.py

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

mlx_lm/benchmark.py

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

mlx_lm/chat.py

@@ -1,7 +1,6 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import argparse
-import json
 
 import mlx.core as mx
 
@@ -28,6 +27,11 @@ def setup_arg_parser():
         help="The path to the local model directory or Hugging Face repo.",
         default=DEFAULT_MODEL,
     )
+    parser.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        help="Enable trusting remote code for tokenizer",
+    )
     parser.add_argument(
         "--adapter-path",
         type=str,
@@ -70,6 +74,11 @@ def setup_arg_parser():
         default=DEFAULT_MAX_TOKENS,
         help="Maximum number of tokens to generate",
     )
+    parser.add_argument(
+        "--system-prompt",
+        default=None,
+        help="System prompt to be used for the chat template",
+    )
     return parser
 
 
@@ -83,7 +92,9 @@ def main():
     model, tokenizer = load(
         args.model,
         adapter_path=args.adapter_path,
-        tokenizer_config={"trust_remote_code": True},
+        tokenizer_config={
+            "trust_remote_code": True if args.trust_remote_code else None
+        },
     )
 
     def print_help():
@@ -105,7 +116,10 @@ def main():
         if query == "h":
             print_help()
             continue
-        messages = [{"role": "user", "content": query}]
+        messages = []
+        if args.system_prompt is not None:
+            messages.append({"role": "system", "content": args.system_prompt})
+        messages.append({"role": "user", "content": query})
         prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
         for response in stream_generate(
             model,

mlx_lm/convert.py

@@ -6,7 +6,7 @@ from typing import Callable, Optional, Union
 
 import mlx.core as mx
 import mlx.nn as nn
-from mlx.utils import tree_flatten
+from mlx.utils import tree_map_with_path
 
 from .utils import (
     dequantize_model,
@@ -19,11 +19,9 @@ from .utils import (
 
 
 def mixed_quant_predicate_builder(
-    recipe: str, model: nn.Module
+    recipe: str, model: nn.Module, group_size: int = 64
 ) -> Callable[[str, nn.Module, dict], Union[bool, dict]]:
-
     high_bits = 6
-    group_size = 64
 
     if recipe == "mixed_2_6":
         low_bits = 2
@@ -35,7 +33,7 @@ def mixed_quant_predicate_builder(
     elif recipe == "mixed_4_6":
         low_bits = 4
     else:
-        raise ValueError("Invalid quant recipe {recipe}")
+        raise ValueError(f"Invalid quant recipe {recipe}")
 
     down_keys = [k for k, _ in model.named_modules() if "down_proj" in k]
     if len(down_keys) == 0:
@@ -50,16 +48,11 @@ def mixed_quant_predicate_builder(
     def mixed_quant_predicate(
         path: str,
         module: nn.Module,
-        config: dict,
     ) -> Union[bool, dict]:
         """Implements mixed quantization predicates with similar choices to, for example, llama.cpp's Q4_K_M.
         Ref: https://github.com/ggerganov/llama.cpp/blob/917786f43d0f29b7c77a0c56767c0fa4df68b1c5/src/llama.cpp#L5265
         By Alex Barron: https://gist.github.com/barronalex/84addb8078be21969f1690c1454855f3
         """
-
-        if not hasattr(module, "to_quantized"):
-            return False
-
         index = (
             int(path.split(".")[layer_location])
             if len(path.split(".")) > layer_location
@@ -70,7 +63,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}
@@ -93,6 +88,7 @@ def convert(
     quantize: bool = False,
     q_group_size: int = 64,
     q_bits: int = 4,
+    q_mode: str = "affine",
     dtype: Optional[str] = None,
     upload_repo: str = None,
     revision: Optional[str] = None,
@@ -100,6 +96,7 @@ def convert(
     quant_predicate: Optional[
         Union[Callable[[str, nn.Module, dict], Union[bool, dict]], str]
     ] = None,
+    trust_remote_code: bool = False,
 ):
     # Check the save path is empty
     if isinstance(mlx_path, str):
@@ -112,47 +109,55 @@ def convert(
         )
 
     print("[INFO] Loading")
-    model_path = get_model_path(hf_path, revision=revision)
-    model, config, tokenizer = fetch_from_hub(model_path, lazy=True)
+    model_path, hf_path = get_model_path(hf_path, revision=revision)
+    model, config, tokenizer = fetch_from_hub(
+        model_path, lazy=True, trust_remote_code=trust_remote_code
+    )
 
     if isinstance(quant_predicate, str):
-        quant_predicate = mixed_quant_predicate_builder(quant_predicate, model)
+        quant_predicate = mixed_quant_predicate_builder(
+            quant_predicate, model, q_group_size
+        )
 
     if dtype is None:
         dtype = config.get("torch_dtype", None)
-    weights = dict(tree_flatten(model.parameters()))
     if dtype in MODEL_CONVERSION_DTYPES:
         print("[INFO] Using dtype:", dtype)
         dtype = getattr(mx, dtype)
+        cast_predicate = getattr(model, "cast_predicate", lambda _: True)
 
-        if hasattr(model, "cast_predicate"):
-            cast_predicate = model.cast_predicate()
-        else:
-            cast_predicate = lambda _: True
-        weights = {
-            k: v.astype(dtype) if cast_predicate(k) else v for k, v in weights.items()
-        }
+        def set_dtype(k, v):
+            if cast_predicate(k) and mx.issubdtype(v.dtype, mx.floating):
+                return v.astype(dtype)
+            else:
+                return v
+
+        model.update(tree_map_with_path(set_dtype, model.parameters()))
 
     if quantize and dequantize:
         raise ValueError("Choose either quantize or dequantize, not both.")
 
     if quantize:
         print("[INFO] Quantizing")
-        model.load_weights(list(weights.items()))
-        weights, config = quantize_model(
-            model, config, q_group_size, q_bits, quant_predicate=quant_predicate
+        model, config = quantize_model(
+            model,
+            config,
+            q_group_size,
+            q_bits,
+            mode=q_mode,
+            quant_predicate=quant_predicate,
         )
 
     if dequantize:
         print("[INFO] Dequantizing")
+        config.pop("quantization", None)
+        config.pop("quantization_config", None)
         model = dequantize_model(model)
-        weights = dict(tree_flatten(model.parameters()))
 
-    del model
     save(
         mlx_path,
         model_path,
-        weights,
+        model,
         tokenizer,
         config,
         hf_repo=hf_path,
@@ -186,6 +191,13 @@ def configure_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--q-bits", help="Bits per weight for quantization.", type=int, default=4
     )
+    parser.add_argument(
+        "--q-mode",
+        help="The quantization mode.",
+        type=str,
+        default="affine",
+        choices=["affine", "mxfp4"],
+    )
     parser.add_argument(
         "--quant-predicate",
         help=f"Mixed-bit quantization recipe.",
@@ -213,6 +225,12 @@ def configure_parser() -> argparse.ArgumentParser:
         action="store_true",
         default=False,
     )
+    parser.add_argument(
+        "--trust-remote-code",
+        help="Trust remote code when loading tokenizer.",
+        action="store_true",
+        default=False,
+    )
     return parser
 
 

mlx_lm/dwq.py

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

mlx_lm/evaluate.py

@@ -20,10 +20,10 @@ import mlx.nn as nn
 import numpy as np
 from lm_eval.api.model import LM
 from lm_eval.api.registry import register_model
+from lm_eval.models import huggingface
 from tqdm import tqdm
 
-from .generate import stream_generate
-from .models.base import create_causal_mask
+from .generate import batch_generate
 from .models.cache import make_prompt_cache
 from .utils import common_prefix_len, load
 
@@ -62,18 +62,22 @@ def chat_template_fn(**extra_kwargs):
 @register_model("mlxlm")
 class MLXLM(LM):
 
-    tokenizer_name = lm_eval.models.huggingface.HFLM.tokenizer_name
+    tokenizer_name = huggingface.HFLM.tokenizer_name
     apply_chat_template = chat_template_fn()
 
     def __init__(
         self,
         path_or_hf_repo: str,
-        max_tokens: Optional[int] = None,
+        max_tokens: int,
         use_chat_template: Optional[bool] = None,
+        trust_remote_code: bool = False,
     ) -> None:
         super().__init__()
-        self._model, self.tokenizer = load(path_or_hf_repo)
-        self._max_tokens = max_tokens or self.tokenizer.model_max_length
+        tokenizer_config = {"trust_remote_code": True if trust_remote_code else None}
+        self._model, self.tokenizer = load(
+            path_or_hf_repo, tokenizer_config=tokenizer_config
+        )
+        self._max_tokens = max_tokens
         self._batch_size = 8
         self.use_chat_template = use_chat_template
         if use_chat_template is None:
@@ -94,30 +98,28 @@ class MLXLM(LM):
         inputs, targets = inputs[..., :-1], inputs[..., 1:]
 
         cache = cache or make_prompt_cache(self._model)
-        lengths += cache[0].offset
-
+        offset = 0
         scores, is_greedy = [], []
         for i in range(0, inputs.shape[1], step_size):
             inp = inputs[:, i : i + step_size]
             T = inp.shape[1]
 
-            offset = cache[0].offset
-            mask = create_causal_mask(T, offset, lengths=lengths)
-
-            logits = self._model(inp, cache=cache, mask=mask)
+            logits = self._model(inp, cache=cache)
             log_probs = nn.log_softmax(logits.astype(mx.float32))
 
             score = mx.take_along_axis(
                 log_probs, targets[:, i : i + step_size, mx.newaxis], axis=-1
             )[..., 0]
+
             ig = targets[:, i : i + step_size] == mx.argmax(logits, axis=-1)
-            ig = mx.where(mx.arange(T) + offset < lengths[:, None], ig, False)
+            ig = mx.where(mx.arange(offset, T + offset) < lengths[:, None], ig, False)
 
             mx.eval(score, ig)
             mx.clear_cache()
 
             is_greedy.append(ig)
             scores.append(score)
+            offset += T
 
         scores = mx.concatenate(scores, axis=1)
         is_greedy = mx.concatenate(is_greedy, axis=1)
@@ -165,7 +167,7 @@ class MLXLM(LM):
         indices = []
         for v in group_reqs.values():
             idx, resp = zip(*v)
-            indices.extend(idx)
+            indices.append(idx)
             responses.append(resp)
 
         # split data accross ranks
@@ -211,31 +213,36 @@ class MLXLM(LM):
                 scores[-1] += mx.sum(score).item()
                 is_greedy[-1] &= mx.all(ig).item()
 
-        scores = mx.array(scores)
-        is_greedy = mx.array(is_greedy)
-
         if long_completions > 0:
             logging.info(
                 f"Prefix eliminated for {long_completions} requests with "
                 + "completion longer than context."
             )
 
+        # All gather the results across nodes
         num_results = len(requests)
+        per_group = mx.distributed.all_max(len(scores), stream=mx.cpu).item()
+        scores = scores + [0] * (per_group - len(scores))
+        is_greedy = is_greedy + [False] * (per_group - len(is_greedy))
+        scores = mx.array(scores)
+        is_greedy = mx.array(is_greedy)
+        scores = mx.distributed.all_gather(scores, stream=mx.cpu)
+        is_greedy = mx.distributed.all_gather(is_greedy, stream=mx.cpu)
+        mx.eval(scores, is_greedy)
 
-        # all gather the results across groups
-        if group.size() > 1:
-            per_group = int(np.ceil(num_results / group.size()))
-            scores = mx.pad(scores, ((0, per_group - len(scores)),))
-            is_greedy = mx.pad(is_greedy, ((0, per_group - len(is_greedy))))
-            scores = mx.distributed.all_gather(scores[mx.newaxis], stream=mx.cpu)
-            is_greedy = mx.distributed.all_gather(is_greedy[mx.newaxis], stream=mx.cpu)
-            mx.eval(scores, is_greedy)
-            scores = scores.T.reshape(-1)
-            is_greedy = is_greedy.T.reshape(-1)
-
-        inv_sort = mx.argsort(mx.array(indices))
+        # Arrange the indices to match the scores from each node and then
+        # inverse sort the scores
+        all_indices = []
+        for rank in range(group.size()):
+            rank_indices = [
+                idx for question in indices[rank :: group.size()] for idx in question
+            ]
+            rank_indices += [num_results] * (per_group - len(rank_indices))
+            all_indices.extend(rank_indices)
+        inv_sort = mx.argsort(mx.array(all_indices))
         scores = scores[:num_results][inv_sort]
         is_greedy = is_greedy[:num_results][inv_sort]
+
         return list(zip(scores.tolist(), is_greedy.tolist()))
 
     def loglikelihood_rolling(self, requests) -> list[float]:
@@ -275,8 +282,8 @@ class MLXLM(LM):
         )
         inputs = self._tokenize([req.args[0] for req in requests])
         all_scores = []
-        for i in tqdm(range(0, len(texts), self._batch_size)):
-            batch = texts[i : i + self._batch_size]
+        for i in tqdm(range(0, len(inputs), self._batch_size)):
+            batch = inputs[i : i + self._batch_size]
             scores, lengths, _ = self._score_fn(batch)
             mask = mx.arange(scores.shape[-1]) < lengths[:, None]
             all_scores.extend((mask * scores).sum(axis=-1).tolist())
@@ -299,30 +306,33 @@ class MLXLM(LM):
         """
         logging.info("Generating continuation for %d sequences." % len(requests))
         contexts, options = zip(*[req.args for req in requests])
-        # contrary to the doc the second element of the tuple contains
+        # The second element of the tuple contains:
         # {'do_sample': False, 'until': ['\n\n'], 'temperature': 0}
         completions = []
 
-        for context, opt in tqdm(zip(contexts, options), total=len(contexts)):
-            until = opt["until"]
-            context = self.tokenizer.encode(
+        # Tokenize all contexts
+        contexts = [
+            self.tokenizer.encode(
                 context, add_special_tokens=not self.use_chat_template
             )
-            max_tokens = min(
-                opt.get("max_gen_tokens", self._max_tokens),
-                self.tokenizer.model_max_length - len(context),
-            )
-            text = ""
-            for response in stream_generate(
-                self._model, self.tokenizer, prompt=context, max_tokens=max_tokens
-            ):
-                text += response.text
-                if any(u in text for u in until):
-                    text = _rstrip_until(text, until)
-                    completions.append(text)
-                    break
-            else:
-                completions.append(text)
+            for context in contexts
+        ]
+
+        # TODO consider multi-token, per-prompt stop conditions
+        max_tokens = [opt.get("max_gen_toks", self._max_tokens) for opt in options]
+
+        completions = batch_generate(
+            model=self._model,
+            tokenizer=self.tokenizer,
+            prompts=contexts,
+            max_tokens=max_tokens,
+            verbose=True,
+        ).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)
         return completions
 
 
@@ -340,7 +350,8 @@ def main():
     parser.add_argument(
         "--max-tokens",
         type=int,
-        help="Maximum nunber of tokens to generate. Defaults to the model's max context length.",
+        help="Maximum number of tokens to generate.",
+        default=8912,
     )
     parser.add_argument(
         "--limit",
@@ -367,10 +378,21 @@ def main():
     parser.add_argument(
         "--chat-template-args",
         type=json.loads,
-        help="""A JSON formatted string of arguments for the tokenizer's "
-        "apply_chat_template, e.g. '{"enable_thinking":false}'""",
+        help="""A JSON formatted string of arguments for the tokenizer's
+        apply_chat_template, e.g. '{"enable_thinking":false}'""",
         default="{}",
     )
+    parser.add_argument(
+        "--confirm-run-unsafe-code",
+        action="store_true",
+        help="Confirm that you want to run tasks that execute untrusted code.",
+        default=False,
+    )
+    parser.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        help="Enable trusting remote code for tokenizer",
+    )
 
     args = parser.parse_args()
 
@@ -382,10 +404,17 @@ def main():
 
     mx.random.seed(args.seed)
 
+    # Initialize the communication if in distributed mode
+    world = mx.distributed.init()
+    mx.eval(mx.distributed.all_sum(1, stream=mx.cpu))
+    if world.size() > 1 and world.rank() == 0:
+        print(f"Evaluating with {world.size()} nodes")
+
     lm = MLXLM(
         args.model,
         max_tokens=args.max_tokens,
         use_chat_template=args.apply_chat_template,
+        trust_remote_code=args.trust_remote_code,
     )
     MLXLM.apply_chat_template = chat_template_fn(**args.chat_template_args)
 
@@ -400,6 +429,7 @@ def main():
         numpy_random_seed=args.seed,
         torch_random_seed=args.seed,
         fewshot_random_seed=args.seed,
+        confirm_run_unsafe_code=args.confirm_run_unsafe_code,
     )
 
     file_keys = ["eval", args.model.replace("/", "_"), version("lm_eval")]
@@ -407,7 +437,7 @@ def main():
         file_keys += [f"{args.num_shots:02d}"]
     file_keys += args.tasks
     filename = "_".join(file_keys)
-    if mx.distributed.init().rank() == 0:
+    if world.rank() == 0:
         output_path = output_dir / filename
         output_path.write_text(json.dumps(results["results"], indent=4))
         print("Results:")

mlx_lm/examples/batch_generate_response.py

@@ -0,0 +1,32 @@
+# Copyright © 2025 Apple Inc.
+
+from mlx_lm import batch_generate, load
+
+# Specify the checkpoint
+checkpoint = "mlx-community/Llama-3.2-3B-Instruct-4bit"
+
+# Load the corresponding model and tokenizer
+model, tokenizer = load(path_or_hf_repo=checkpoint)
+
+# A batch of prompts
+prompts = [
+    "Write a story about Einstein.",
+    "Why is the sky blue?",
+    "What time is it?",
+    "How tall is Mt Everest?",
+]
+
+# Apply the chat template and encode to tokens
+prompts = [
+    tokenizer.apply_chat_template(
+        [{"role": "user", "content": p}],
+        add_generation_prompt=True,
+    )
+    for p in prompts
+]
+
+# Set `verbose=True` to see generation statistics
+result = batch_generate(model, tokenizer, prompts, verbose=False, max_tokens=128)
+
+# The returned result contains texts completions in the same order as prompts
+print(result.texts[0])

mlx_lm/examples/lora_config.yaml

@@ -1,5 +1,5 @@
 # The path to the local model directory or Hugging Face repo.
-model: "mlx_model"
+model: "mlx-community/Llama-3.2-1B-Instruct-bf16"
 
 # Whether or not to train (boolean)
 train: true
@@ -17,7 +17,7 @@ optimizer: adamw
 #     bias_correction: true
 
 # Directory with {train, valid, test}.jsonl files
-data: "/path/to/training/data"
+data: "mlx-community/WikiSQL"
 
 # The PRNG seed
 seed: 0
@@ -37,6 +37,10 @@ val_batches: 25
 # Adam learning rate.
 learning_rate: 1e-5
 
+# Services to report logs to (comma-separated): wandb, swanlab, or both ('wandb,swanlab').
+# report_to: wandb,swanlab
+# project_name: "Your-awesome-mlx-project-name"
+
 # Number of training steps between loss reporting.
 steps_per_report: 10
 

mlx_lm/examples/openai_tool_use.py

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

mlx_lm/examples/pipeline_generate.py

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

mlx_lm/fuse.py

@@ -4,8 +4,6 @@ from pathlib import Path
 from mlx.utils import tree_flatten, tree_unflatten
 
 from .gguf import convert_to_gguf
-from .tuner.dora import DoRAEmbedding, DoRALinear
-from .tuner.lora import LoRAEmbedding, LoRALinear, LoRASwitchLinear
 from .tuner.utils import dequantize, load_adapters
 from .utils import (
     fetch_from_hub,
@@ -35,12 +33,6 @@ def parse_arguments() -> argparse.Namespace:
         default="adapters",
         help="Path to the trained adapter weights and config.",
     )
-    parser.add_argument(
-        "--hf-path",
-        type=str,
-        default=None,
-        help="Path to the original Hugging Face model. Required for upload if --model is a local directory.",
-    )
     parser.add_argument(
         "--upload-repo",
         help="The Hugging Face repo to upload the model to.",
@@ -70,14 +62,16 @@ def main() -> None:
     print("Loading pretrained model")
     args = parse_arguments()
 
-    model_path = get_model_path(args.model)
+    model_path, hf_path = get_model_path(args.model)
     model, config, tokenizer = fetch_from_hub(model_path)
 
     model.freeze()
     model = load_adapters(model, args.adapter_path)
 
     fused_linears = [
-        (n, m.fuse()) for n, m in model.named_modules() if hasattr(m, "fuse")
+        (n, m.fuse(de_quantize=args.de_quantize))
+        for n, m in model.named_modules()
+        if hasattr(m, "fuse")
     ]
 
     if fused_linears:
@@ -88,18 +82,15 @@ def main() -> None:
         model = dequantize(model)
         config.pop("quantization", None)
 
-    weights = dict(tree_flatten(model.parameters()))
-
     save_path = Path(args.save_path)
-    hf_path = args.hf_path or (args.model if not Path(args.model).exists() else None)
     save(
         save_path,
         model_path,
-        weights,
+        model,
         tokenizer,
         config,
         hf_repo=hf_path,
-        donate_weights=False,
+        donate_model=False,
     )
 
     if args.export_gguf:
@@ -108,6 +99,7 @@ def main() -> None:
             raise ValueError(
                 f"Model type {model_type} not supported for GGUF conversion."
             )
+        weights = dict(tree_flatten(model.parameters()))
         convert_to_gguf(model_path, weights, config, str(save_path / args.gguf_path))
 
     if args.upload_repo is not None:
@@ -115,7 +107,7 @@ def main() -> None:
             raise ValueError(
                 "Must provide original Hugging Face repo to upload local model."
             )
-        upload_to_hub(args.save_path, args.upload_repo, hf_path)
+        upload_to_hub(args.save_path, args.upload_repo)
 
 
 if __name__ == "__main__":

mlx_lm/generate.py

@@ -24,14 +24,18 @@ from transformers import PreTrainedTokenizer
 
 from .models import cache
 from .models.cache import (
+    ArraysCache,
+    BatchKVCache,
+    BatchRotatingKVCache,
+    CacheList,
+    KVCache,
     QuantizedKVCache,
+    RotatingKVCache,
     load_prompt_cache,
-    make_prompt_cache,
-    trim_prompt_cache,
 )
 from .sample_utils import make_sampler
 from .tokenizer_utils import TokenizerWrapper
-from .utils import load
+from .utils import does_model_support_input_embeddings, load
 
 DEFAULT_PROMPT = "hello"
 DEFAULT_MAX_TOKENS = 100
@@ -63,6 +67,11 @@ def setup_arg_parser():
         ),
         default=None,
     )
+    parser.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        help="Enable trusting remote code for tokenizer",
+    )
     parser.add_argument(
         "--adapter-path",
         type=str,
@@ -216,29 +225,35 @@ def wired_limit(model: nn.Module, streams: Optional[List[mx.Stream]] = None):
     async eval could be running pass in the streams to synchronize with prior
     to exiting the context manager.
     """
-    model_bytes = tree_reduce(
-        lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, model, 0
-    )
-    max_rec_size = mx.metal.device_info()["max_recommended_working_set_size"]
-    if model_bytes > 0.9 * max_rec_size:
-        model_mb = model_bytes // 2**20
-        max_rec_mb = max_rec_size // 2**20
-        print(
-            f"[WARNING] Generating with a model that requires {model_mb} MB "
-            f"which is close to the maximum recommended size of {max_rec_mb} "
-            "MB. This can be slow. See the documentation for possible work-arounds: "
-            "https://github.com/ml-explore/mlx-lm/tree/main#large-models"
+    if not mx.metal.is_available():
+        try:
+            yield
+        finally:
+            pass
+    else:
+        model_bytes = tree_reduce(
+            lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, model, 0
         )
-    old_limit = mx.set_wired_limit(max_rec_size)
-    try:
-        yield None
-    finally:
-        if streams is not None:
-            for s in streams:
-                mx.synchronize(s)
-        else:
-            mx.synchronize()
-        mx.set_wired_limit(old_limit)
+        max_rec_size = mx.metal.device_info()["max_recommended_working_set_size"]
+        if model_bytes > 0.9 * max_rec_size:
+            model_mb = model_bytes // 2**20
+            max_rec_mb = max_rec_size // 2**20
+            print(
+                f"[WARNING] Generating with a model that requires {model_mb} MB "
+                f"which is close to the maximum recommended size of {max_rec_mb} "
+                "MB. This can be slow. See the documentation for possible work-arounds: "
+                "https://github.com/ml-explore/mlx-lm/tree/main#large-models"
+            )
+        old_limit = mx.set_wired_limit(max_rec_size)
+        try:
+            yield
+        finally:
+            if streams is not None:
+                for s in streams:
+                    mx.synchronize(s)
+            else:
+                mx.synchronize()
+            mx.set_wired_limit(old_limit)
 
 
 @dataclass
@@ -272,16 +287,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(
@@ -298,6 +308,7 @@ def generate_step(
     kv_group_size: int = 64,
     quantized_kv_start: int = 0,
     prompt_progress_callback: Optional[Callable[int, int]] = None,
+    input_embeddings: Optional[mx.array] = None,
 ) -> Generator[Tuple[mx.array, mx.array], None, None]:
     """
     A generator producing token ids based on the given prompt from the model.
@@ -322,14 +333,28 @@ def generate_step(
         kv_group_size (int): Group size for KV cache quantization. Default: ``64``.
         quantized_kv_start (int): Step to begin using a quantized KV cache.
            when ``kv_bits`` is non-None. Default: ``0``.
-        prompt_prorgress_callback (Callable[int, int]): A call-back which takes the
+        prompt_progress_callback (Callable[int, int]): A call-back which takes the
            prompt tokens processed so far and the total number of prompt tokens.
+        input_embeddings (mx.array, optional): Input embeddings to use instead of or in
+          conjunction with prompt tokens. Default: ``None``.
 
     Yields:
         Tuple[mx.array, mx.array]: One token and a vector of log probabilities.
     """
+    if input_embeddings is not None:
+        if not does_model_support_input_embeddings(model):
+            raise ValueError("Model does not support input embeddings.")
+        elif len(prompt) > 0 and len(prompt) != len(input_embeddings):
+            raise ValueError(
+                f"When providing input_embeddings, their sequence length ({len(input_embeddings)}) "
+                f"must match the sequence length of the prompt ({len(prompt)}), or the "
+                "prompt must be empty."
+            )
+    elif len(prompt) == 0:
+        raise ValueError(
+            "Either input_embeddings or prompt (or both) must be provided."
+        )
 
-    y = prompt
     tokens = None
 
     # Create the KV cache for generation
@@ -338,8 +363,6 @@ def generate_step(
             model,
             max_kv_size=max_kv_size,
         )
-    elif len(prompt_cache) != len(model.layers):
-        raise ValueError("Wrong number of layers in the prompt cache.")
 
     prompt_progress_callback = prompt_progress_callback or (lambda *_: None)
 
@@ -352,37 +375,71 @@ def generate_step(
 
     sampler = sampler or (lambda x: mx.argmax(x, axis=-1))
 
-    def _step(y):
+    def _model_call(input_tokens: mx.array, input_embeddings: Optional[mx.array]):
+        if input_embeddings is not None:
+            return model(
+                input_tokens, cache=prompt_cache, input_embeddings=input_embeddings
+            )
+        else:
+            return model(input_tokens, cache=prompt_cache)
+
+    def _step(input_tokens: mx.array, input_embeddings: Optional[mx.array] = None):
+        nonlocal tokens
+
         with mx.stream(generation_stream):
-            logits = model(y[None], cache=prompt_cache)
+            logits = _model_call(
+                input_tokens=input_tokens[None],
+                input_embeddings=(
+                    input_embeddings[None] if input_embeddings is not None else None
+                ),
+            )
+
             logits = logits[:, -1, :]
 
-            if logits_processors:
-                nonlocal tokens
-                tokens = mx.concat([tokens, y]) if tokens is not None else y
-
+            if logits_processors and len(input_tokens) > 0:
+                tokens = (
+                    mx.concat([tokens, input_tokens])
+                    if tokens is not None
+                    else input_tokens
+                )
                 for processor in logits_processors:
                     logits = processor(tokens, logits)
 
             quantize_cache_fn(prompt_cache)
 
             logprobs = logits - mx.logsumexp(logits, keepdims=True)
-            y = sampler(logprobs)
-            return y, logprobs.squeeze(0)
+            sampled = sampler(logprobs)
+            return sampled, logprobs.squeeze(0)
 
     with mx.stream(generation_stream):
-        total_prompt_tokens = y.size
+        total_prompt_tokens = (
+            len(input_embeddings) if input_embeddings is not None else len(prompt)
+        )
         prompt_processed_tokens = 0
-        while y.size > prefill_step_size:
-            model(y[:prefill_step_size][None], cache=prompt_cache)
+        prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens)
+        while total_prompt_tokens - prompt_processed_tokens > 1:
+            n_to_process = min(prefill_step_size, prompt.size - 1)
+            _model_call(
+                input_tokens=prompt[:n_to_process][None],
+                input_embeddings=(
+                    input_embeddings[:n_to_process][None]
+                    if input_embeddings is not None
+                    else None
+                ),
+            )
             quantize_cache_fn(prompt_cache)
             mx.eval([c.state for c in prompt_cache])
+            prompt_processed_tokens += n_to_process
             prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens)
-            prompt_processed_tokens += prefill_step_size
-            y = y[prefill_step_size:]
+            prompt = prompt[n_to_process:]
+            input_embeddings = (
+                input_embeddings[n_to_process:]
+                if input_embeddings is not None
+                else input_embeddings
+            )
             mx.clear_cache()
 
-        y, logprobs = _step(y)
+        y, logprobs = _step(input_tokens=prompt, input_embeddings=input_embeddings)
 
     mx.async_eval(y, logprobs)
     n = 0
@@ -407,7 +464,7 @@ def speculative_generate_step(
     model: nn.Module,
     draft_model: nn.Module,
     *,
-    num_draft_tokens=2,
+    num_draft_tokens: int = 2,
     max_tokens: int = 256,
     sampler: Optional[Callable[mx.array, mx.array]] = None,
     logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None,
@@ -454,8 +511,6 @@ def speculative_generate_step(
     if prompt_cache is None:
         model_cache = cache.make_prompt_cache(model)
         draft_cache = cache.make_prompt_cache(draft_model)
-    elif len(prompt_cache) != (len(model.layers) + len(draft_model.layers)):
-        raise ValueError("Wrong number of layers in the prompt cache.")
     else:
         model_cache = prompt_cache[: len(model.layers)]
         draft_cache = prompt_cache[len(model.layers) :]
@@ -583,6 +638,7 @@ def stream_generate(
     model: nn.Module,
     tokenizer: Union[PreTrainedTokenizer, TokenizerWrapper],
     prompt: Union[str, mx.array, List[int]],
+    max_tokens: int = 256,
     draft_model: Optional[nn.Module] = None,
     **kwargs,
 ) -> Generator[GenerationResponse, None, None]:
@@ -594,6 +650,8 @@ def stream_generate(
         tokenizer (PreTrainedTokenizer): The tokenizer.
         prompt (Union[str, mx.array, List[int]]): The input prompt string or
           integer tokens.
+        max_tokens (int): The maximum number of tokens to generate.
+          Default: ``256``.
         draft_model (Optional[nn.Module]): An optional draft model. If provided
           then speculative decoding is used. The draft model must use the same
           tokenizer as the main model. Default: ``None``.
@@ -618,6 +676,8 @@ def stream_generate(
 
     detokenizer = tokenizer.detokenizer
 
+    kwargs["max_tokens"] = max_tokens
+
     if draft_model is None:
         kwargs.pop("num_draft_tokens", None)
         token_generator = generate_step(prompt, model, **kwargs)
@@ -627,11 +687,11 @@ def stream_generate(
         )
     else:
         kwargs.pop("max_kv_size", None)
+        kwargs.pop("prompt_progress_callback", None)
         token_generator = speculative_generate_step(
             prompt, model, draft_model, **kwargs
         )
     with wired_limit(model, [generation_stream]):
-        detokenizer.reset()
         tic = time.perf_counter()
         for n, (token, logprobs, from_draft) in enumerate(token_generator):
             if n == 0:
@@ -642,6 +702,8 @@ def stream_generate(
                 break
 
             detokenizer.add_token(token)
+            if (n + 1) == max_tokens:
+                break
 
             yield GenerationResponse(
                 text=detokenizer.last_segment,
@@ -676,7 +738,6 @@ def generate(
     tokenizer: Union[PreTrainedTokenizer, TokenizerWrapper],
     prompt: Union[str, List[int]],
     verbose: bool = False,
-    formatter: Optional[Callable] = None,
     **kwargs,
 ) -> str:
     """
@@ -691,11 +752,6 @@ def generate(
        kwargs: The remaining options get passed to :func:`stream_generate`.
           See :func:`stream_generate` for more details.
     """
-    if formatter is not None:
-        print(
-            "[Warning] Text formatting is deprecated and no longer used. "
-            "The argument will be removed in a future version."
-        )
     if verbose:
         print("=" * 10)
 
@@ -723,6 +779,340 @@ def generate(
     return text
 
 
+def _left_pad_prompts(prompts, max_length=None):
+    if max_length is None:
+        max_length = max(len(p) for p in prompts)
+    return mx.array([[0] * (max_length - len(p)) + p for p in prompts])
+
+
+@dataclass
+class BatchStats:
+    """
+    An data object to hold generation stats.
+
+    Args:
+        prompt_tokens (int): The number of prompt tokens processed.
+        prompt_tps (float): The prompt processing tokens-per-second.
+        prompt_time (float): The time in seconds spent in prompt processing.
+        generation_tokens (int): The number of generated tokens.
+        generation_tps (float): The tokens-per-second for generation.
+        generation_time (float): The time in seconds spent in generation .
+        peak_memory (float): The peak memory used so far in GB.
+    """
+
+    prompt_tokens: int = 0
+    prompt_tps: float = 0
+    prompt_time: float = 0
+    generation_tokens: int = 0
+    generation_tps: float = 0
+    generation_time: float = 0
+    peak_memory: float = 0
+
+
+@dataclass
+class BatchResponse:
+    """
+    An data object to hold a batch generation response.
+
+    Args:
+        texts: (List[str]): The generated text for each prompt.
+        stats (BatchStats): Statistics about the generation.
+    """
+
+    texts: List[str]
+    stats: BatchStats
+
+
+@dataclass
+class Batch:
+    uids: List[int]
+    y: mx.array
+    logprobs: mx.array
+    max_tokens: List[int]
+    num_tokens: List[int]
+    cache: List[Any]
+
+    def __len__(self):
+        return len(self.uids)
+
+    def filter(self, keep_idx: List[int]):
+        self.uids = [self.uids[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.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 _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()
+        return [to_batch_cache(c) for c in cache]
+    else:
+        return [BatchKVCache(left_padding) for _ in model.layers]
+
+
+class BatchGenerator:
+
+    @dataclass
+    class Response:
+        uid: int
+        token: int
+        logprobs: mx.array
+        finish_reason: Optional[str]
+
+    def __init__(
+        self,
+        model,
+        max_tokens: int = 128,
+        stop_tokens: Optional[set] = None,
+        sampler: Optional[Callable[mx.array, mx.array]] = None,
+        completion_batch_size: int = 32,
+        prefill_batch_size: int = 8,
+        prefill_step_size: int = 2048,
+    ):
+        self.model = model
+        self.unprocessed_prompts = []
+        self.max_tokens = max_tokens
+        self.stop_tokens = stop_tokens or set()
+        self.sampler = sampler or (lambda x: mx.argmax(x, axis=-1))
+        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._stats = BatchStats()
+
+        self.active_batch = None
+
+    def insert(self, prompts, max_tokens: Union[List[int], int, None] = 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))
+            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])
+        )
+        return uids
+
+    def _process_prompts(self, prompts):
+        uids, inputs, max_tokens = zip(*prompts)
+        lengths = [len(p) for p in inputs]
+        max_length = max(lengths)
+        batch_size = self.prefill_batch_size
+        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)
+
+        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:]
+            mx.clear_cache()
+
+        y, logprobs = self._step(inputs, prompt_cache)
+        mx.async_eval(y, logprobs)
+        return Batch(
+            list(uids), y, logprobs, list(max_tokens), [0] * len(uids), prompt_cache
+        )
+
+    def _step(self, input_tokens: mx.array, prompt_cache: List[Any]):
+        logits = self.model(input_tokens, cache=prompt_cache)
+        logits = logits[:, -1, :]
+        logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
+        sampled = self.sampler(logprobs)
+        return sampled, logprobs
+
+    def stats(self):
+        self._stats.prompt_tps = self._stats.prompt_tokens / self._stats.prompt_time
+        self._stats.generation_tps = (
+            self._stats.generation_tokens / self._stats.generation_time
+        )
+        self._stats.peak_memory = mx.get_peak_memory() / 1e9
+        return self._stats
+
+    def _next(self):
+        tic = time.perf_counter()
+
+        prompt_processing = False
+        batch = self.active_batch
+        num_active = len(batch) if batch else 0
+        num_to_add = self.completion_batch_size - num_active
+        while num_to_add >= self.prefill_batch_size:
+            prompts = self.unprocessed_prompts[: self.prefill_batch_size]
+            # Finish processing the last examples of the last batch
+            if len(prompts) == 0 and num_active > 0:
+                break
+            # No more prompts and no more completions, all done
+            elif len(prompts) == 0:
+                self.active_batch = None
+                return []
+            # Process prompts
+            if batch is not None and not prompt_processing:
+                # Finish any active completion tokens
+                mx.eval(batch.y, batch.logprobs)
+                self._stats.generation_time += time.perf_counter() - tic
+                tic = time.perf_counter()
+
+            batch = self._process_prompts(prompts)
+            self.unprocessed_prompts = self.unprocessed_prompts[
+                self.prefill_batch_size :
+            ]
+            prompt_processing = True
+            # If there was no active batch, set it
+            if self.active_batch is None:
+                self.active_batch = batch
+            else:
+                self.active_batch.extend(batch)
+
+            num_active = len(self.active_batch)
+            num_to_add -= len(batch)
+
+        batch = self.active_batch
+        y, logprobs = batch.y, batch.logprobs
+        batch.y, batch.logprobs = self._step(y[:, None], batch.cache)
+        mx.async_eval(batch.y, batch.logprobs)
+
+        y = y.tolist()
+        toc = time.perf_counter()
+        if prompt_processing:
+            self._stats.prompt_time += toc - tic
+        else:
+            self._stats.generation_time += toc - tic
+        keep_idx = []
+        end_idx = []
+        responses = []
+
+        for e, (t, uid, num_tok, max_tok) in enumerate(
+            zip(y, batch.uids, batch.num_tokens, batch.max_tokens)
+        ):
+            num_tok += 1
+            batch.num_tokens[e] = num_tok
+            if t in self.stop_tokens:
+                finish_reason = "stop"
+                end_idx.append(e)
+            elif num_tok >= max_tok:
+                finish_reason = "length"
+                end_idx.append(e)
+            else:
+                finish_reason = None
+                keep_idx.append(e)
+            responses.append(self.Response(uid, t, logprobs[e], finish_reason))
+
+        # Remove any finished completions
+        if len(end_idx):
+            if len(keep_idx) > 0:
+                batch.filter(keep_idx)
+            else:
+                self.active_batch = None
+
+        self._stats.generation_tokens += len(responses)
+        return responses
+
+    def next(self):
+        with mx.stream(generation_stream):
+            return self._next()
+
+
+def batch_generate(
+    model,
+    tokenizer,
+    prompts: List[int],
+    max_tokens: Union[int, List[int]] = 128,
+    verbose: bool = False,
+    **kwargs,
+) -> BatchResponse:
+    """
+    Generate responses for the given batch of prompts.
+
+    Args:
+       model (nn.Module): The language model.
+       tokenizer (PreTrainedTokenizer): The tokenizer.
+       prompt (List[List[int]]): The input prompts.
+       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.
+       kwargs: The remaining options get passed to :obj:`BatchGenerator`.
+          See :obj:`BatchGenerator` for more details.
+    """
+
+    gen = BatchGenerator(model, stop_tokens=tokenizer.eos_token_ids, **kwargs)
+    num_samples = len(prompts)
+    fin = 0
+    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:
+                    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 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()
+    if verbose:
+        print(
+            f"[batch_generate] Prompt: {stats.prompt_tokens} tokens, {stats.prompt_tps:.3f} tokens-per-sec"
+        )
+        print(
+            f"[batch_generate] Generation: {stats.generation_tokens} tokens, "
+            f"{stats.generation_tps:.3f} tokens-per-sec"
+        )
+        print(f"[batch_generate] Peak memory: {stats.peak_memory:.3f} GB")
+    return BatchResponse(texts, stats)
+
+
 def main():
     parser = setup_arg_parser()
     args = parser.parse_args()
@@ -751,7 +1141,7 @@ def main():
     tokenizer_config = (
         {} if not using_cache else json.loads(metadata["tokenizer_config"])
     )
-    tokenizer_config["trust_remote_code"] = True
+    tokenizer_config["trust_remote_code"] = True if args.trust_remote_code else None
 
     model_path = args.model
     if using_cache:

mlx_lm/lora.py

@@ -1,10 +1,9 @@
-# Copyright © 2024 Apple Inc.
-
 import argparse
 import math
 import os
 import re
 import types
+import warnings
 from pathlib import Path
 
 import mlx.core as mx
@@ -13,6 +12,7 @@ import mlx.optimizers as optim
 import numpy as np
 import yaml
 
+from .tuner.callbacks import get_reporting_callbacks
 from .tuner.datasets import CacheDataset, load_dataset
 from .tuner.trainer import TrainingArgs, TrainingCallback, evaluate, train
 from .tuner.utils import (
@@ -40,15 +40,18 @@ yaml_loader.add_implicit_resolver(
 )
 
 CONFIG_DEFAULTS = {
-    "model": "mlx_model",
+    "model": "Qwen/Qwen3-0.6b",
     "train": False,
     "fine_tune_type": "lora",
     "optimizer": "adam",
     "optimizer_config": {
         "adam": {},
         "adamw": {},
+        "muon": {},
+        "sgd": {},
+        "adafactor": {},
     },
-    "data": "data/",
+    "data": "mlx-community/WikiSQL",
     "seed": 0,
     "num_layers": 16,
     "batch_size": 4,
@@ -66,8 +69,11 @@ CONFIG_DEFAULTS = {
     "config": None,
     "grad_checkpoint": False,
     "lr_schedule": None,
-    "lora_parameters": {"rank": 8, "dropout": 0.0, "scale": 10.0},
+    "lora_parameters": {"rank": 8, "dropout": 0.0, "scale": 20.0},
     "mask_prompt": False,
+    "wandb": None,  # will be deprecated in a future release
+    "report_to": None,
+    "project_name": None,
 }
 
 
@@ -103,9 +109,9 @@ def build_parser():
     parser.add_argument(
         "--optimizer",
         type=str,
-        choices=["adam", "adamw"],
+        choices=["adam", "adamw", "muon", "sgd", "adafactor"],
         default=None,
-        help="Optimizer to use for training: adam or adamw",
+        help="Optimizer to use for training: adam, adamw, sgd, or adafactor.",
     )
     parser.add_argument(
         "--mask-prompt",
@@ -179,6 +185,27 @@ def build_parser():
         help="Use gradient checkpointing to reduce memory use.",
         default=None,
     )
+    parser.add_argument(  # will be deprecated in a future release
+        "--wandb",
+        type=str,
+        default=None,
+        help=(
+            "The 'wandb' argument is deprecated and will be removed in a future release. "
+            "Use 'report_to: wandb' and 'project_name' in the configuration instead."
+        ),
+    )
+    parser.add_argument(
+        "--report-to",
+        type=str,
+        default=None,
+        help="Services to report logs to ('wandb', 'swanlab', or 'wandb,swanlab').",
+    )
+    parser.add_argument(
+        "--project-name",
+        type=str,
+        default=None,
+        help="Project name for logging. Defaults to the name of the root directory.",
+    )
     parser.add_argument("--seed", type=int, help="The PRNG seed")
     return parser
 
@@ -201,6 +228,8 @@ def train_model(
     if args.fine_tune_type == "full":
         for l in model.layers[-max(args.num_layers, 0) :]:
             l.unfreeze()
+
+        args.lora_parameters = None
     elif args.fine_tune_type in ["lora", "dora"]:
         # Convert linear layers to lora/dora layers and unfreeze in the process
         linear_to_lora_layers(
@@ -243,11 +272,16 @@ def train_model(
 
     optimizer_name = args.optimizer.lower()
     optimizer_config = args.optimizer_config.get(optimizer_name, {})
-
     if optimizer_name == "adam":
         opt_class = optim.Adam
     elif optimizer_name == "adamw":
         opt_class = optim.AdamW
+    elif optimizer_name == "muon":
+        opt_class = optim.Muon
+    elif optimizer_name == "sgd":
+        opt_class = optim.SGD
+    elif optimizer_name == "adafactor":
+        opt_class = optim.Adafactor
     else:
         raise ValueError(f"Unsupported optimizer: {optimizer_name}")
 
@@ -280,9 +314,23 @@ 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,
+        log_dir=args.adapter_path,
+        config=vars(args),
+    )
 
     print("Loading pretrained model")
-    model, tokenizer = load(args.model)
+    model, tokenizer = load(args.model, tokenizer_config={"trust_remote_code": True})
 
     print("Loading datasets")
     train_set, valid_set, test_set = load_dataset(args, tokenizer)

mlx_lm/merge.py

@@ -1,176 +0,0 @@
-# Copyright © 2023-2024 Apple Inc.
-
-import argparse
-import glob
-import shutil
-from pathlib import Path
-from typing import Optional
-
-import mlx.core as mx
-import mlx.nn as nn
-import numpy as np
-import yaml
-from mlx.utils import tree_flatten, tree_map
-
-from .utils import (
-    fetch_from_hub,
-    get_model_path,
-    save_config,
-    save_weights,
-    upload_to_hub,
-)
-
-
-def configure_parser() -> argparse.ArgumentParser:
-    """
-    Configures and returns the argument parser for the script.
-
-    Returns:
-        argparse.ArgumentParser: Configured argument parser.
-    """
-    parser = argparse.ArgumentParser(description="Merge multiple models.")
-
-    parser.add_argument("--config", type=str, help="Path to the YAML config.")
-    parser.add_argument(
-        "--mlx-path",
-        type=str,
-        default="mlx_merged_model",
-        help="Path to save the MLX model.",
-    )
-    parser.add_argument(
-        "--upload-repo",
-        help="The Hugging Face repo to upload the model to.",
-        type=str,
-        default=None,
-    )
-    return parser
-
-
-def slerp(t, w1, w2, eps=1e-5):
-    """
-    Spherical linear interpolation
-
-    Args:
-        t (float): Interpolation weight in [0.0, 1.0]
-        w1 (mx.array): First input
-        w2 (mx.array): Second input
-        eps (float): Constant for numerical stability
-    Returns:
-        mx.array: Interpolated result
-    """
-    t = float(t)
-    if t == 0:
-        return w1
-    elif t == 1:
-        return w2
-    # Normalize
-    v1 = w1 / mx.linalg.norm(w1)
-    v2 = w2 / mx.linalg.norm(w2)
-    # Angle
-    dot = mx.clip((v1 * v2).sum(), 0.0, 1.0)
-    theta = mx.arccos(dot)
-    sin_theta = mx.sin(theta + eps)
-    s1 = mx.sin(theta * (1 - t)) / sin_theta
-    s2 = mx.sin(theta * t) / sin_theta
-    return s1 * w1 + s2 * w2
-
-
-def merge_models(base_model: nn.Module, model: nn.Module, config: dict):
-    method = config.get("method", None)
-    if method != "slerp":
-        raise ValueError(f"Merge method {method} not supported")
-
-    num_layers = len(model.layers)
-
-    def unpack_values(vals):
-        if isinstance(vals, (int, float)):
-            return np.full(num_layers, vals)
-        bins = len(vals) - 1
-        sizes = [num_layers // bins] * bins
-        sizes[-1] = num_layers - sum(sizes[:-1])
-        return np.concatenate(
-            [np.linspace(v1, v2, s) for v1, v2, s in zip(vals[:-1], vals[1:], sizes)]
-        )
-
-    param_list = config["parameters"]["t"]
-    params = {}
-    filter_keys = set()
-    for pl in param_list[:-1]:
-        params[pl["filter"]] = unpack_values(pl["value"])
-        filter_keys.add(pl["filter"])
-    default = unpack_values(param_list[-1]["value"])
-
-    for e in range(num_layers):
-        bl = base_model.layers[e]
-        l = model.layers[e]
-        base_weights = bl.parameters()
-        weights = l.parameters()
-        for k, w1 in base_weights.items():
-            w2 = weights[k]
-            t = params.get(k, default)[e]
-            base_weights[k] = tree_map(lambda x, y: slerp(t, x, y), w1, w2)
-        base_model.update(base_weights)
-
-
-def merge(
-    config: str,
-    mlx_path: str = "mlx_model",
-    upload_repo: Optional[str] = None,
-):
-    with open(config, "r") as fid:
-        merge_conf = yaml.safe_load(fid)
-    print("[INFO] Loading")
-
-    model_paths = merge_conf.get("models", [])
-    if len(model_paths) < 2:
-        raise ValueError(f"Expected at least 2 models, got {len(model_paths)}.")
-
-    # Load all models
-    base_hf_path = model_paths[0]
-    base_path = get_model_path(base_hf_path)
-    base_model, base_config, tokenizer = fetch_from_hub(base_path, lazy=True)
-    models = []
-    for mp in model_paths[1:]:
-        model, model_config, _ = fetch_from_hub(get_model_path(mp), lazy=True)
-        base_type = base_config["model_type"]
-        model_type = model_config["model_type"]
-        if base_type != model_type:
-            raise ValueError(
-                f"Can only merge models of the same type,"
-                f" but got {base_type} and {model_type}."
-            )
-        models.append(model)
-
-    # Merge models into base model
-    for m in models:
-        merge_models(base_model, m, merge_conf)
-
-    # Save base model
-    mlx_path = Path(mlx_path)
-    weights = dict(tree_flatten(base_model.parameters()))
-    del models, base_model
-    save_weights(mlx_path, weights, donate_weights=True)
-    py_files = glob.glob(str(base_path / "*.py"))
-    for file in py_files:
-        shutil.copy(file, mlx_path)
-
-    tokenizer.save_pretrained(mlx_path)
-
-    save_config(config, config_path=mlx_path / "config.json")
-
-    if upload_repo is not None:
-        upload_to_hub(mlx_path, upload_repo, base_hf_path)
-
-
-def main():
-    parser = configure_parser()
-    args = parser.parse_args()
-    merge(**vars(args))
-
-
-if __name__ == "__main__":
-    print(
-        "Calling `python -m mlx_lm.merge...` directly is deprecated."
-        " Use `mlx_lm.merge...` or `python -m mlx_lm merge ...` instead."
-    )
-    main()

mlx_lm/models/Klear.py

@@ -0,0 +1,262 @@
+# 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
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    attention_bias: bool
+    mlp_only_layers: List[int]
+    num_experts: int
+    num_experts_per_tok: int
+    decoder_sparse_step: int
+    n_shared_experts: int
+    moe_intermediate_size: int
+    rms_norm_eps: float
+    vocab_size: int
+    num_key_value_heads: int
+    rope_theta: float
+    max_position_embeddings: int
+    norm_topk_prob: bool
+
+
+class KlearAttention(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,
+            self.num_attention_heads * self.head_dim,
+            bias=args.attention_bias,
+        )
+        self.k_proj = nn.Linear(
+            args.hidden_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=args.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            args.hidden_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=args.attention_bias,
+        )
+        self.o_proj = nn.Linear(
+            self.num_attention_heads * self.head_dim,
+            args.hidden_size,
+            bias=args.attention_bias,
+        )
+
+        self.q_norm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
+        self.k_norm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
+
+        self.rope = nn.RoPE(
+            self.head_dim,
+            traditional=False,
+            base=args.rope_theta,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        queries = self.q_norm(
+            queries.reshape(B, L, self.num_attention_heads, -1)
+        ).transpose(0, 2, 1, 3)
+        keys = self.k_norm(keys.reshape(B, L, self.num_key_value_heads, -1)).transpose(
+            0, 2, 1, 3
+        )
+        values = values.reshape(B, L, self.num_key_value_heads, -1).transpose(
+            0, 2, 1, 3
+        )
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class KlearMLP(nn.Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class KlearSparseMoeBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.norm_topk_prob = args.norm_topk_prob
+        self.num_experts = args.num_experts
+        self.top_k = args.num_experts_per_tok
+
+        self.gate = nn.Linear(args.hidden_size, args.num_experts, bias=False)
+        self.experts = SwitchGLU(
+            args.hidden_size, args.moe_intermediate_size, args.num_experts
+        )
+        self.shared_experts = KlearMLP(
+            args.hidden_size,
+            hidden_dim=args.moe_intermediate_size * args.n_shared_experts,
+        )
+        self.coefficient = nn.Linear(args.hidden_size, 2)
+        self.expert_bias = mx.zeros((self.num_experts,), dtype=mx.float32)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        routing_weights = mx.sigmoid(self.gate(x).astype(mx.float32))
+        biased_weights = routing_weights + self.expert_bias.reshape((1, 1, -1))
+        k = self.top_k
+        inds = mx.argpartition(-biased_weights, kth=k - 1, axis=-1)[..., :k]
+        scores = mx.take_along_axis(routing_weights, inds, axis=-1)
+        if self.norm_topk_prob:
+            scores = scores / mx.sum(scores, axis=-1, keepdims=True)
+        scores = scores.astype(x.dtype)
+        expert_out = self.experts(x, inds)
+        y_experts = (expert_out * scores[..., None]).sum(axis=-2)
+        coef = mx.softmax(self.coefficient(x), axis=-1, precise=True)
+        shared = self.shared_experts(x)
+        y = y_experts * coef[..., :1] + shared * coef[..., 1:]
+        return y
+
+
+class KlearDecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.self_attn = KlearAttention(args)
+
+        if (layer_idx not in args.mlp_only_layers) and (
+            args.num_experts > 0 and (layer_idx + 1) % args.decoder_sparse_step == 0
+        ):
+            self.mlp = KlearSparseMoeBlock(args)
+        else:
+            self.mlp = KlearMLP(args.hidden_size, args.intermediate_size)
+
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        out = h + r
+        return out
+
+
+class KlearModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            KlearDecoderLayer(args=args, layer_idx=i)
+            for i in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        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 = KlearModel(args)
+        self.lm_head = nn.Linear(args.hidden_size, args.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):
+        if "model.layers.0.mlp.experts.0.gate_proj.weight" not in weights:
+            return weights
+
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}.mlp.experts"
+            for name in ["gate_proj", "up_proj", "down_proj"]:
+                stacked = [
+                    weights.pop(f"{prefix}.{e}.{name}.weight")
+                    for e in range(self.args.num_experts)
+                ]
+                weights[f"{prefix}.{name}.weight"] = mx.stack(stacked)
+
+        return weights
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if path.endswith("mlp.gate"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate
+
+    @property
+    def cast_predicate(self):
+        def predicate(k):
+            return "expert_bias" not in k
+
+        return predicate

mlx_lm/models/afm7.py

@@ -0,0 +1,394 @@
+# Copyright © 2025 Apple Inc.
+
+import math
+from dataclasses import dataclass
+from functools import partial
+from itertools import accumulate
+from typing import Any, Dict, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import ConcatenateKVCache, KVCache
+from .rope_utils import initialize_rope
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_dim: int
+    num_layers: int
+    num_kv_reuse_layers: int
+    num_heads: int
+    num_kv_heads: int
+    hidden_dim_scale_factor: float = 3.25
+    rope_theta: float = 50000
+    rms_norm_eps: float = 1e-5
+
+
+class FusedLoRALinear(nn.Module):
+    def __init__(
+        self,
+        input_dims: int,
+        output_dims: list[int],
+        r: int = 8,
+        dropout: float = 0.0,
+        scale: float = 20.0,
+    ):
+        super().__init__()
+
+        self.linear = FusedLinear(input_dims, output_dims)
+        self.dropout = nn.Dropout(p=dropout)
+        self.scale = scale
+
+        scale = 1 / math.sqrt(input_dims)
+        self.lora_a = [
+            mx.random.uniform(low=-scale, high=scale, shape=(input_dims, r))
+            for _ in output_dims
+        ]
+        self.lora_b = [mx.zeros((r, od)) for od in output_dims]
+
+    def fuse(self, de_quantize: bool = False):
+        linear = self.linear
+        weight = linear.weight
+        is_quantized = isinstance(linear, FusedQuantizedLinear)
+
+        # Use the same type as the linear weight if not quantized
+        dtype = weight.dtype
+
+        if is_quantized:
+            dtype = linear.scales.dtype
+            weight = mx.dequantize(
+                weight,
+                linear.scales,
+                linear.biases,
+                linear.group_size,
+                linear.bits,
+            )
+
+        input_dims = weight.shape[-1]
+        output_dims = linear.output_dims
+        fused_linear = FusedLinear(input_dims, output_dims)
+        fused_linear.weight = weight
+        deltas = [
+            ((self.scale * b.T) @ a.T).astype(dtype)
+            for a, b in zip(self.lora_a, self.lora_b)
+        ]
+        delta = mx.concatenate(deltas, axis=0)
+        fused_linear.weight = weight + delta
+
+        if is_quantized and not de_quantize:
+            fused_linear = fused_linear.to_quantized(linear.group_size, linear.bits)
+
+        return fused_linear
+
+    def __call__(self, x):
+        dt = x.dtype
+        y = self.linear(x)
+        x = self.dropout(x)
+        z = [(x @ a) @ b for a, b in zip(self.lora_a, self.lora_b)]
+        return tuple(yi + (self.scale * zi).astype(dt) for yi, zi in zip(y, z))
+
+
+class FusedQuantizedLinear(nn.QuantizedLinear):
+    def __init__(self, input_dims, output_dims, group_size: int = 64, bits: int = 4):
+        *indices, output_dims = accumulate(output_dims)
+        self.indices = indices
+        super().__init__(
+            input_dims, output_dims, bias=False, group_size=group_size, bits=bits
+        )
+
+    @property
+    def input_dims(self):
+        return self.scales.shape[-1] * self.group_size
+
+    @property
+    def output_dims(self):
+        indices = [0] + self.indices + [self.weight.shape[0]]
+        return [indices[i] - indices[i - 1] for i in range(1, len(indices))]
+
+    def __call__(self, x):
+        x = super().__call__(x)
+        return x.split(self.indices, axis=-1)
+
+    def to_lora(self, r: int = 8, dropout: float = 0.0, scale: float = 20.0):
+        lora_lin = FusedLoRALinear(self.input_dims, self.output_dims, r, dropout, scale)
+        lora_lin.linear = self
+        return lora_lin
+
+
+class FusedLinear(nn.Linear):
+    def __init__(self, input_dims, output_dims):
+        *indices, output_dims = accumulate(output_dims)
+        self.indices = indices
+        super().__init__(input_dims, output_dims, bias=False)
+
+    @property
+    def input_dims(self):
+        return self.weight.shape[-1]
+
+    @property
+    def output_dims(self):
+        indices = [0] + self.indices + [self.weight.shape[0]]
+        return [indices[i] - indices[i - 1] for i in range(1, len(indices))]
+
+    def __call__(self, x):
+        x = super().__call__(x)
+        return x.split(self.indices, axis=-1)
+
+    def to_quantized(self, group_size: int = 64, bits: int = 4):
+        input_dims = self.input_dims
+        output_dims = self.output_dims
+        ql = FusedQuantizedLinear(input_dims, output_dims, group_size, bits)
+        ql.weight, ql.scales, ql.biases = mx.quantize(self.weight, group_size, bits)
+
+        return ql
+
+    def to_lora(self, r: int = 8, dropout: float = 0.0, scale: float = 20.0):
+        lora_lin = FusedLoRALinear(self.input_dims, self.output_dims, r, dropout, scale)
+        lora_lin.linear = self
+        return lora_lin
+
+
+@partial(mx.compile, shapeless=True)
+def fake_8bit_quant(x, scale):
+    dt = x.dtype
+    x = x.astype(mx.float32)
+    x = (x / scale).round()
+    x = mx.clip(x, -128, 127)
+    return (x * scale).astype(dt)
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_dim
+        self.n_heads = n_heads = args.num_heads
+        self.n_kv_heads = n_kv_heads = args.num_kv_heads
+        self.head_dim = head_dim = args.hidden_dim // n_heads
+        self.scale = head_dim**-0.5
+
+        qkv_dim = (n_heads + 2 * n_kv_heads) * head_dim
+        self.qkv_proj = FusedLinear(
+            dim, [n_heads * head_dim] + 2 * [n_kv_heads * head_dim]
+        )
+        self.out_proj = nn.Linear(dim, dim, bias=False)
+        self.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            True,
+        )
+        self.q_norm = nn.RMSNorm(head_dim)
+        self.k_norm = nn.RMSNorm(head_dim)
+        self.quant_key_scale = mx.array(1.0)
+        self.quant_value_scale = mx.array(1.0)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        # Get the queries, keys and values
+        queries, keys, values = self.qkv_proj(x)
+
+        # Prepare the queries, keys and values for the attention computation
+        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            queries = self.q_norm(self.rope(queries, offset=cache.offset))
+            keys = self.k_norm(self.rope(keys, offset=cache.offset))
+            keys = fake_8bit_quant(keys, self.quant_key_scale)
+            values = fake_8bit_quant(values, self.quant_value_scale)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.q_norm(self.rope(queries))
+            keys = self.k_norm(self.rope(keys))
+            keys = fake_8bit_quant(keys, self.quant_key_scale)
+            values = fake_8bit_quant(values, self.quant_value_scale)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.out_proj(output)
+
+
+class KVReuseAttention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_dim
+        self.n_heads = n_heads = args.num_heads
+        self.head_dim = head_dim = args.hidden_dim // n_heads
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, dim, bias=False)
+        self.out_proj = nn.Linear(dim, dim, bias=False)
+        self.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            True,
+        )
+        self.q_norm = nn.RMSNorm(head_dim)
+
+    def __call__(
+        self,
+        x: mx.array,
+        keys: mx.array,
+        values: mx.array,
+        mask: Optional[mx.array] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+        _, _, S, _ = keys.shape
+
+        queries = self.q_proj(x)
+        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        queries = self.q_norm(self.rope(queries, offset=S - L))
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=None, scale=self.scale, mask=mask
+        )
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.out_proj(output)
+
+
+@partial(mx.compile, shapeless=True)
+def _swiglu(g, x):
+    return nn.silu(g) * x
+
+
+class MLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_dim
+        hidden_dim = int(dim * args.hidden_dim_scale_factor)
+
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
+
+    def __call__(self, x) -> mx.array:
+        g = self.gate_proj(x)
+        x = self.up_proj(x)
+        return self.down_proj(_swiglu(g, x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.self_attn = Attention(args)
+        self.mlp = MLP(args)
+        self.input_layernorm = nn.RMSNorm(args.hidden_dim, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_dim, eps=args.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        out = h + r
+        return out
+
+
+class KVReuseTransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.self_attn = KVReuseAttention(args)
+        self.mlp = MLP(args)
+        self.input_layernorm = nn.RMSNorm(args.hidden_dim, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_dim, eps=args.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        keys: mx.array,
+        values: mx.array,
+        mask: Optional[mx.array] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), keys, values, mask)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        out = h + r
+        return out
+
+
+class AFMModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+
+        self.embedding = nn.Embedding(args.vocab_size, args.hidden_dim)
+        self.layers = [
+            TransformerBlock(args)
+            for _ in range(args.num_layers - args.num_kv_reuse_layers)
+        ]
+        self.kv_reuse_layers = [
+            KVReuseTransformerBlock(args) for _ in range(args.num_kv_reuse_layers)
+        ]
+        self.output_norm = nn.RMSNorm(args.hidden_dim, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+    ):
+        h = self.embedding(inputs)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+            cache[-1] = ConcatenateKVCache()
+
+        mask = create_attention_mask(h, cache[0])
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, cache=c)
+
+        keys, values = cache[-1].state
+        for layer in self.kv_reuse_layers:
+            h = layer(h, keys, values, mask)
+
+        return self.output_norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = AFMModel(args)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+    ):
+        out = self.model(inputs, cache)
+        out = self.model.embedding.as_linear(out)
+        return out
+
+    def make_cache(self):
+        return [KVCache() for _ in range(len(self.model.layers))]
+
+    @property
+    def layers(self):
+        return self.model.layers + self.model.kv_reuse_layers

mlx_lm/models/apertus.py

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

mlx_lm/models/baichuan_m1.py

@@ -96,7 +96,10 @@ class Attention(nn.Module):
         k = k.reshape(B, L, self.num_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
         v = v.reshape(B, L, self.num_kv_heads, self.head_dim).transpose(0, 2, 1, 3)
 
-        if cache is not None:
+        if cache is None:
+            cache = (None, None)
+
+        if cache[0] is not None:
             offset = cache[1].offset
             last_k, last_v = cache[0][0], cache[0][1]
         else:
@@ -110,7 +113,7 @@ class Attention(nn.Module):
         q = self.rope(q, offset=offset)
         k = self.rope(k, offset=offset)
 
-        if cache is not None:
+        if cache[0] is not None:
             k, v = cache[1].update_and_fetch(k, v)
             if L > 0:
                 cache[0][0] = k_init[:, :, -1:, :]
@@ -167,17 +170,40 @@ class BaichuanModel(nn.Module):
         self.layers = [DecoderLayer(config, i) for i in range(config.num_hidden_layers)]
         self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
-    def __call__(
-        self, inputs: mx.array, mask: mx.array = None, cache: Any = None
-    ) -> mx.array:
+        self.sliding_window = config.sliding_window
+        self.first_swa_idx = None
+        if config.sliding_window_layers:
+            self.first_swa_idx = config.sliding_window_layers[0]
+
+        self.first_global_idx = None
+        self.swa_layers = set(config.sliding_window_layers)
+        for i in range(config.num_hidden_layers):
+            if i in self.swa_layers:
+                continue
+            self.first_global_idx = i
+            break
+
+    def __call__(self, inputs: mx.array, cache: Any = None) -> mx.array:
         x = self.embed_tokens(inputs)
-        if mask is None:
-            if cache is not None:
-                c = [cache[0][1]]
-            mask = create_attention_mask(x, c)
+
         if cache is None:
-            cache = [None] * len(self.layers)
-        for layer, c in zip(self.layers, cache):
+            cache = [(None, None)] * len(self.layers)
+
+        if self.first_global_idx is None:
+            c_global = None
+        else:
+            c_global = cache[self.first_global_idx][1]
+
+        if self.first_swa_idx is None:
+            c_swa = None
+        else:
+            c_swa = cache[self.first_swa_idx][1]
+
+        global_mask = create_attention_mask(x, c_global)
+        swa_mask = create_attention_mask(x, c_swa, window_size=self.sliding_window)
+
+        for l, (layer, c) in enumerate(zip(self.layers, cache)):
+            mask = swa_mask if l in self.swa_layers else global_mask
             x = layer(x, mask, c)
         return self.norm(x)
 
@@ -215,10 +241,8 @@ class Model(nn.Module):
             weights["lm_head.weight"] = w
         return weights
 
-    def __call__(
-        self, inputs: mx.array, mask: mx.array = None, cache: Any = None
-    ) -> mx.array:
-        outputs = self.model(inputs, mask, cache)
+    def __call__(self, inputs: mx.array, cache: Any = None) -> mx.array:
+        outputs = self.model(inputs, cache)
         return self.lm_head(outputs)
 
     @property

mlx_lm/models/bailing_moe.py

@@ -0,0 +1,391 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    intermediate_size: int
+    max_position_embeddings: int
+    moe_intermediate_size: int
+    num_experts: int
+    num_shared_experts: int
+    norm_topk_prob: bool
+    num_attention_heads: int
+    num_experts_per_tok: int
+    num_hidden_layers: int
+    num_key_value_heads: int
+    rms_norm_eps: float
+    rope_theta: float
+    vocab_size: int
+    first_k_dense_replace: int
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    use_bias: bool = False
+    use_qkv_bias: bool = False
+    norm_head: bool = False
+    norm_softmax: bool = False
+    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
+    score_function: str = "softmax"
+    n_group: int = 1
+    topk_group: int = 4
+    moe_shared_expert_intermediate_size: Optional[int] = None
+    moe_router_enable_shared_expert: bool = True
+
+
+class BailingMoeMLP(nn.Module):
+    def __init__(self, args: ModelArgs, intermediate_size: Optional[int] = None):
+        super().__init__()
+        self.intermediate_size = (
+            intermediate_size
+            if intermediate_size is not None
+            else args.intermediate_size
+        )
+
+        self.gate_proj = nn.Linear(
+            args.hidden_size, self.intermediate_size, bias=args.use_bias
+        )
+        self.down_proj = nn.Linear(
+            self.intermediate_size, args.hidden_size, bias=args.use_bias
+        )
+        self.up_proj = nn.Linear(
+            args.hidden_size, self.intermediate_size, bias=args.use_bias
+        )
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class BailingMoeAttention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.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.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)
+
+        if (rope_dim := args.rotary_dim) is None:
+            rope_dim = int(self.head_dim * args.partial_rotary_factor)
+        self.rope = initialize_rope(
+            rope_dim,
+            args.rope_theta,
+            traditional=False,
+            scaling_config=args.rope_scaling,
+            max_position_embeddings=args.max_position_embeddings,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        qkv = self.query_key_value(x)
+
+        q_size = self.num_attention_heads * self.head_dim
+        kv_size = self.num_key_value_heads * self.head_dim
+        q, k, v = mx.split(qkv, [q_size, q_size + kv_size], axis=-1)
+
+        queries = q.reshape(B, L, self.num_attention_heads, self.head_dim).transpose(
+            0, 2, 1, 3
+        )
+        keys = k.reshape(B, L, self.num_key_value_heads, self.head_dim).transpose(
+            0, 2, 1, 3
+        )
+        values = v.reshape(B, L, self.num_key_value_heads, self.head_dim).transpose(
+            0, 2, 1, 3
+        )
+
+        if 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)
+
+
+def group_expert_select(
+    gates,
+    e_score_correction_bias,
+    top_k,
+    n_group,
+    topk_group,
+    routed_scaling_factor,
+    norm_topk_prob,
+    score_function,
+):
+
+    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 BailingMoeGate(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):
+        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 BailingMoeSparseMoeBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.num_experts_per_tok = args.num_experts_per_tok
+        self.switch_mlp = SwitchGLU(
+            args.hidden_size,
+            args.moe_intermediate_size,
+            args.num_experts,
+            bias=args.use_bias,
+        )
+        self.gate = BailingMoeGate(args)
+        shared_dim = (
+            args.moe_shared_expert_intermediate_size or args.moe_intermediate_size
+        )
+        self.shared_experts = (
+            BailingMoeMLP(
+                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):
+        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 BailingMoeDecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.attention = BailingMoeAttention(args)
+
+        self.mlp = (
+            BailingMoeSparseMoeBlock(args)
+            if (
+                args.num_experts is not None and layer_idx >= args.first_k_dense_replace
+            )
+            else BailingMoeMLP(args)
+        )
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.attention(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        return h + r
+
+
+class BailingMoeModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.word_embeddings = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            BailingMoeDecoderLayer(args, layer_idx=i)
+            for i in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ):
+        h = self.word_embeddings(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.norm_head = args.norm_head
+        self.model_type = args.model_type
+        self.model = BailingMoeModel(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.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}"
+
+            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

mlx_lm/models/base.py

@@ -7,8 +7,6 @@ from typing import Any, Optional
 import mlx.core as mx
 from mlx.utils import tree_map
 
-from .cache import QuantizedKVCache
-
 
 @dataclass
 class BaseModelArgs:
@@ -27,7 +25,8 @@ def create_causal_mask(
     N: int,
     offset: int = 0,
     window_size: Optional[int] = None,
-    lengths: Optional[mx.array] = None,
+    right_padding: Optional[mx.array] = None,
+    left_padding: Optional[mx.array] = None,
 ):
     rinds = mx.arange(offset + N)
     linds = mx.arange(offset, offset + N) if offset else rinds
@@ -35,34 +34,31 @@ def create_causal_mask(
     rinds = rinds[None]
     mask = linds >= rinds
     if window_size is not None:
-        mask = mask & (linds <= rinds + window_size)
-    if lengths is not None:
-        lengths = lengths[:, None, None, None]
-        mask = mask & (rinds < lengths)
+        mask = mask & (linds < rinds + window_size)
+    if right_padding is not None:
+        mask = mask & (rinds < mx.expand_dims((offset + N) - right_padding, (1, 2, 3)))
+    if left_padding is not None:
+        mask = mask & (mx.expand_dims(left_padding, (1, 2, 3)) <= rinds)
     return mask
 
 
 def create_attention_mask(
-    h: mx.array, cache: Optional[Any] = None, return_array: bool = False
+    h, cache=None, window_size: Optional[int] = None, return_array: bool = False
 ):
-    T = h.shape[1]
-    if T > 1:
-        offset = 0
-        window_size = None
-        if cache is not None and cache[0] is not None:
-            c = cache[0]
-            offset = c.offset
-            if hasattr(c, "max_size"):
-                window_size = c.max_size
-                offset = min(window_size, offset)
-                return_array = return_array or offset + T > window_size
-        if return_array:
-            return create_causal_mask(T, offset, window_size=window_size)
-        else:
-            return "causal"
-    else:
-        mask = None
-    return mask
+    N = h.shape[1]
+    if cache and hasattr(cache, "make_mask"):
+        return cache.make_mask(N, return_array=return_array, window_size=window_size)
+    if N == 1:
+        return None
+    if return_array or (window_size and N > window_size):
+        return create_causal_mask(N, window_size=window_size)
+    return "causal"
+
+
+def create_ssm_mask(h, cache=None):
+    if cache and hasattr(cache, "make_mask"):
+        return cache.make_mask(h.shape[1])
+    return None
 
 
 def quantized_scaled_dot_product_attention(
@@ -116,8 +112,11 @@ def scaled_dot_product_attention(
     cache,
     scale: float,
     mask: Optional[mx.array],
+    sinks: Optional[mx.array] = None,
 ) -> mx.array:
-    if isinstance(cache, QuantizedKVCache):
+    if hasattr(cache, "bits"):
+        if sinks is not None:
+            raise ValueError("Quantized SDPA does not support attention sinks.")
         return quantized_scaled_dot_product_attention(
             queries,
             keys,
@@ -129,5 +128,10 @@ def scaled_dot_product_attention(
         )
     else:
         return mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=scale, mask=mask
+            queries,
+            keys,
+            values,
+            scale=scale,
+            mask=mask,
+            sinks=sinks,
         )

mlx_lm/models/bitlinear_layers.py

@@ -0,0 +1,158 @@
+# Copyright © 2025 Apple Inc.
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.nn.layers.quantized import QuantizedLinear
+from mlx.utils import tree_flatten, tree_unflatten
+
+
+def bitnet_quantize(model, quantization_config: dict):
+    quantize_layers = []
+    modules_to_not_convert = quantization_config.get("modules_to_not_convert", [])
+    invert_weight_scales = (
+        quantization_config.get("linear_class", "") != "autobitlinear"
+    )
+
+    for name, module in tree_flatten(model.leaf_modules(), is_leaf=nn.Module.is_module):
+
+        # Replace nn.Linear layers, but skip any layer from the `modules_to_not_convert` list
+        if name not in modules_to_not_convert and isinstance(module, nn.Linear):
+            old_weight = module.weight
+            out_features, in_features = old_weight.shape
+            bias = "bias" in module
+            new_layer = BitLinear(
+                in_features,
+                out_features,
+                bias=bias,
+                invert_weight_scales=invert_weight_scales,
+            )
+            quantize_layers.append((name, new_layer))
+    if len(quantize_layers) > 0:
+        model.update_modules(tree_unflatten(quantize_layers))
+    return model
+
+
+def make_bitlinear_kernel():
+    """
+    Custom Metal kernel that performs matrix multiplication directly on
+    packed weights and scales the output. This eliminates the need to
+    store unpacked weights in memory.
+    """
+    source = """
+    constexpr int M = 4;
+    constexpr int BLOCK = 32;
+
+    uint tid = thread_position_in_grid.y;
+    uint in_offset = thread_position_in_grid.x;
+
+    uint batch_idx = tid / (out_features / 4);
+    uint row_idx = tid % (out_features / 4);
+
+    float sum[4] = {0.0};
+
+    for (uint i = in_offset * M; i < in_features; i += BLOCK * M) {
+        float v[M];
+        for (int j=0; j<M; j++) {
+            v[j] = x[batch_idx * in_features + i + j];
+        }
+
+        for (int j=0; j<M; j++) {
+            uint8_t w = packed_weights[row_idx * in_features + i + j];
+            sum[0] += v[j] * ((w & 3) - 1);
+            sum[1] += v[j] * (((w >> 2) & 3) - 1);
+            sum[2] += v[j] * (((w >> 4) & 3) - 1);
+            sum[3] += v[j] * (((w >> 6) & 3) - 1);
+        }
+    }
+
+    for (int j=0; j<4; j++) {
+        sum[j] = simd_sum(sum[j]);
+    }
+
+    // Apply weight scaling by diving them or multiplying them
+    if (in_offset == 0) {
+        float scale = invert_weight_scales ? 1 / weight_scale[0] : weight_scale[0];
+        for (int i=0; i<4; i++) {
+            out[batch_idx * out_features + row_idx + i * (out_features/4)] = static_cast<T>(sum[i] * scale);
+        }
+    }
+    """
+
+    return mx.fast.metal_kernel(
+        name="bitlinear_matmul",
+        input_names=["x", "packed_weights", "weight_scale"],
+        output_names=["out"],
+        source=source,
+    )
+
+
+_bitlinear_kernel = make_bitlinear_kernel()
+
+
+class BitLinear(nn.Module):
+    """
+    BitLinear module with memory-efficient weight handling.
+    """
+
+    def __init__(
+        self,
+        in_features,
+        out_features,
+        bias=True,
+        invert_weight_scales=False,
+    ):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        # Calculate packed dimensions - the first dimension gets packed 4:1
+        # The weights are ternary so can be represented with 2 bits, and they
+        # are packed in uint8 tensors, hence the number of values per item is 4
+        packed_out_features = (out_features + 3) // 4
+        self.weight = mx.zeros((packed_out_features, in_features), dtype=mx.uint8)
+
+        self.invert_weight_scales = invert_weight_scales
+        self.weight_scale = mx.array([1.0])
+
+        if bias:
+            self.bias = mx.zeros((out_features,))
+        else:
+            self.bias = None
+
+    def execute_matmul_kernel(self, x, packed_weights):
+        original_shape = x.shape
+        if len(original_shape) > 2:
+            x = x.reshape(-1, original_shape[-1])
+        total_batch_elements, in_features = x.shape
+
+        out_features = self.out_features
+
+        dtype = self.weight_scale.dtype
+        assert x.dtype == dtype, "Wrong type for input."
+        out = _bitlinear_kernel(
+            inputs=[
+                x,
+                packed_weights,
+                self.weight_scale,
+            ],
+            template=[
+                ("T", dtype),
+                ("invert_weight_scales", self.invert_weight_scales),
+                ("in_features", in_features),
+                ("out_features", out_features),
+            ],
+            grid=(32, total_batch_elements * out_features // 4, 1),
+            threadgroup=(32, 1, 1),  # SIMD width is 32 threads
+            output_shapes=[(total_batch_elements, out_features)],
+            output_dtypes=[dtype],
+        )[0]
+
+        if len(original_shape) > 2:
+            out = out.reshape(*original_shape[:-1], out_features)
+        return out
+
+    def __call__(self, x):
+        y = self.execute_matmul_kernel(x, self.weight)
+
+        if self.bias is not None:
+            y = mx.add(y, self.bias)
+        return y

mlx_lm/models/bitnet.py

@@ -0,0 +1,208 @@
+# Copyright © 2023-2024 Apple Inc.
+
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, Dict, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .bitlinear_layers import BitLinear
+from .rope_utils import initialize_rope
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    rms_norm_eps: float
+    vocab_size: int
+    head_dim: Optional[int] = None
+    max_position_embeddings: Optional[int] = None
+    attention_bias: bool = False
+    mlp_bias: bool = False
+    rope_theta: float = 10000
+    rope_traditional: bool = False
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = True
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+
+        self.head_dim = head_dim = args.head_dim or args.hidden_size // n_heads
+
+        self.scale = head_dim**-0.5
+        attention_bias = args.attention_bias
+
+        self.q_proj = BitLinear(dim, n_heads * head_dim, bias=attention_bias)
+        self.k_proj = BitLinear(dim, n_kv_heads * head_dim, bias=attention_bias)
+        self.v_proj = BitLinear(dim, n_kv_heads * head_dim, bias=attention_bias)
+        self.o_proj = BitLinear(n_heads * head_dim, dim, bias=attention_bias)
+
+        self.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            args.rope_traditional,
+            args.rope_scaling,
+            args.max_position_embeddings,
+        )
+        self.attn_sub_norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        # Prepare the queries, keys and values for the attention computation
+        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        output = self.attn_sub_norm(output)
+        output = self.o_proj(output)
+
+        return output
+
+
+class MLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        hidden_dim = args.intermediate_size
+        if hasattr(args, "mlp_bias"):
+            mlp_bias = args.mlp_bias
+        else:
+            mlp_bias = False
+
+        self.gate_proj = BitLinear(dim, hidden_dim, bias=mlp_bias)
+        self.down_proj = BitLinear(hidden_dim, dim, bias=mlp_bias)
+        self.up_proj = BitLinear(dim, hidden_dim, bias=mlp_bias)
+
+        self.ffn_sub_norm = nn.RMSNorm(args.intermediate_size, eps=args.rms_norm_eps)
+
+    def __call__(self, x) -> mx.array:
+        x = nn.relu2(self.gate_proj(x)) * self.up_proj(x)
+        x = self.ffn_sub_norm(x)
+        x = self.down_proj(x)
+        return x
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.hidden_size = args.hidden_size
+        self.self_attn = Attention(args)
+        self.mlp = MLP(args)
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        out = h + r
+        return out
+
+
+class LlamaModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.num_hidden_layers = args.num_hidden_layers
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            TransformerBlock(args=args) for _ in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+    ):
+        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, cache=c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = LlamaModel(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        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 "self_attn.rotary_emb.inv_freq" not in k
+        }
+        if self.args.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
+        return weights
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/cache.py

@@ -1,18 +1,21 @@
 # Copyright © 2023-2024 Apple Inc.
 
+import copy
 from typing import Any, Dict, List, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
 from mlx.utils import tree_flatten, tree_map, tree_unflatten
 
+from .base import create_causal_mask
+
 
 def make_prompt_cache(
     model: nn.Module,
     max_kv_size: Optional[int] = None,
 ) -> List[Any]:
     """
-    Construct the model's cache for use when cgeneration.
+    Construct the model's cache for use in generation.
 
     This function will defer the cache construction to the model if it has a
     ``make_cache`` method, otherwise it will make a default KV cache.
@@ -71,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
@@ -106,6 +109,17 @@ def trim_prompt_cache(cache: List[Any], num_tokens: int) -> List[Any]:
     return [c.trim(num_tokens) for c in cache][0]
 
 
+def create_attention_mask(
+    N: int, offset: int, return_array: bool, window_size: Optional[int]
+):
+    if N == 1:
+        return None
+    if return_array:
+        return create_causal_mask(N, offset, window_size=window_size)
+    else:
+        return "causal"
+
+
 class _BaseCache:
     @property
     def state(self):
@@ -128,13 +142,67 @@ class _BaseCache:
     def is_trimmable(self):
         return False
 
+    @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.
+
+    Can be used as a mock KV cache or when large blocks are being processed at
+    a time in which case KVCache isn't necessarily faster. Consider using the
+    KVCache with a larger step size before using this cache.
+    """
+
+    def __init__(self):
+        self.keys = None
+        self.values = None
+        self.offset = 0
+
+    def update_and_fetch(self, keys, values):
+        if self.keys is None:
+            self.keys = keys
+            self.values = values
+        else:
+            self.keys = mx.concatenate([self.keys, keys], axis=-2)
+            self.values = mx.concatenate([self.values, values], axis=-2)
+        self.offset = self.keys.shape[-2]
+
+        return self.keys, self.values
+
+    @property
+    def state(self):
+        return self.keys, self.values
+
+    @state.setter
+    def state(self, v):
+        self.keys, self.values = v
+        self.offset = self.keys.shape[-2]
+
+    def is_trimmable(self):
+        return True
+
+    def trim(self, n):
+        n = min(self.offset, n)
+        self.offset -= n
+        return n
+
+    def make_mask(self, *args, **kwargs):
+        return create_attention_mask(*args, offset=self.offset, **kwargs)
+
 
 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
 
@@ -196,11 +264,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
@@ -210,13 +278,17 @@ class QuantizedKVCache(_BaseCache):
         self.offset -= n
         return n
 
+    def make_mask(self, *args, **kwargs):
+        return create_attention_mask(*args, offset=self.offset, **kwargs)
+
 
 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
@@ -275,16 +347,19 @@ class KVCache(_BaseCache):
             )
         return quant_cache
 
+    def make_mask(self, *args, **kwargs):
+        return create_attention_mask(*args, offset=self.offset, **kwargs)
+
 
 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):
@@ -324,10 +399,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]
@@ -395,13 +471,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,
         )
@@ -418,10 +492,37 @@ class RotatingKVCache(_BaseCache):
     def to_quantized(self, group_size: int = 64, bits: int = 4) -> QuantizedKVCache:
         raise NotImplementedError("RotatingKVCache Quantization NYI")
 
+    def make_mask(
+        self, N: int, window_size: Optional[int] = None, return_array: bool = False
+    ):
+        if N > 1:
+            window_size = window_size or self.max_size
+            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:
+                return "causal"
+        else:
+            if window_size is None:
+                return None
+            # May need a mask for when window_size < max_size
+            if self.offset >= window_size and self.max_size > window_size:
+                idx = self._idx
+                if idx >= self.max_size:
+                    idx = 0
+                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
 
-class MambaCache(_BaseCache):
-    def __init__(self):
-        self.cache = [None, None]
+
+class ArraysCache(_BaseCache):
+    def __init__(self, size, left_padding: Optional[List[int]] = None):
+        self.cache = [None] * size
+        self.left_padding = mx.array(left_padding) if left_padding else None
 
     def __setitem__(self, idx, value):
         self.cache[idx] = value
@@ -437,9 +538,34 @@ class MambaCache(_BaseCache):
     def state(self, v):
         self.cache = v
 
+    def filter(self, batch_indices):
+        """
+        In-place filter to keep just the given indices in the cache.
+        """
+        self.cache = [c[batch_indices] for c in self.cache]
+        self.left_padding = None
+
+    def extend(self, other):
+        """
+        In-place extend this cache with the other cache.
+        """
+        self.cache = [mx.concatenate([c, o]) for c, o in zip(self.cache, other.cache)]
+        self.left_padding = None
+
+    def make_mask(self, N: int):
+        if self.cache[0] is None and self.left_padding is not None:
+            return mx.arange(N) >= self.left_padding[:, None]
+        else:
+            return None
+
+
+class MambaCache(ArraysCache):
+    def __init__(self, left_padding: Optional[List[int]] = None):
+        super().__init__(size=2, left_padding=left_padding)
+
 
 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
@@ -490,13 +616,21 @@ 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
 
     def __getitem__(self, idx):
         return self.caches[idx]
 
+    def is_trimmable(self):
+        return all(c.is_trimmable() for c in self.caches)
+
+    def trim(self, n):
+        for c in self.caches:
+            m = c.trim(n)
+        return m
+
     @property
     def state(self):
         return [s for c in self.caches for s in c.state]
@@ -509,3 +643,365 @@ class CacheList(KVCache):
             l = len(c.state)
             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)
+
+
+class BatchKVCache(_BaseCache):
+    step = 256
+
+    def __init__(self, left_padding: List[int]):
+        """
+        The BatchKV cache expects inputs to be left-padded.
+
+        E.g. the following prompts:
+
+            [1, 3, 5]
+            [7]
+            [2, 6, 8, 9]
+
+        Should be padded like so:
+
+            [0, 1, 3, 5]
+            [0, 0, 0, 7]
+            [2, 6, 8, 9]
+
+        And ``left_padding`` specifies the amount of padding for each.
+        In this case, ``left_padding = [1, 3, 0]``.
+        """
+        self.keys = None
+        self.values = None
+        self.left_padding = mx.array(left_padding)
+        self.offset = mx.array([-l for l in left_padding])
+        self._idx = 0
+
+    def update_and_fetch(self, keys, values):
+        prev = self._idx
+        if self.keys is None or (prev + keys.shape[2]) > self.keys.shape[2]:
+            B, n_kv_heads, _, k_head_dim = keys.shape
+            v_head_dim = values.shape[3]
+            n_steps = (self.step + keys.shape[2] - 1) // self.step
+            k_shape = (B, n_kv_heads, n_steps * self.step, k_head_dim)
+            v_shape = (B, n_kv_heads, n_steps * self.step, v_head_dim)
+            new_k = mx.zeros(k_shape, keys.dtype)
+            new_v = mx.zeros(v_shape, values.dtype)
+            if self.keys is not None:
+                if prev % self.step != 0:
+                    self.keys = self.keys[..., :prev, :]
+                    self.values = self.values[..., :prev, :]
+                self.keys = mx.concatenate([self.keys, new_k], axis=2)
+                self.values = mx.concatenate([self.values, new_v], axis=2)
+            else:
+                self.keys, self.values = new_k, new_v
+
+        self.offset += keys.shape[2]
+        self._idx += keys.shape[2]
+        self.keys[..., prev : self._idx, :] = keys
+        self.values[..., prev : self._idx, :] = values
+        return self.keys[..., : self._idx, :], self.values[..., : self._idx, :]
+
+    @property
+    def state(self):
+        k, v = self.keys, self.values
+        if self._idx < k.shape[2]:
+            k = k[..., : self._idx, :]
+            v = v[..., : self._idx, :]
+        return k, v, self.offset, self.left_padding
+
+    @state.setter
+    def state(self, v):
+        self.keys, self.values, self.offset, self.left_padding = v
+        self._idx = self.keys.shape[2]
+
+    def is_trimmable(self):
+        return True
+
+    def trim(self, n):
+        n = min(self._idx, n)
+        self._idx -= n
+        self.offset -= n
+        return n
+
+    def make_mask(self, N: int, return_array: bool = False, **kwargs):
+        return create_causal_mask(
+            N, offset=self._idx, left_padding=self.left_padding, **kwargs
+        )
+
+    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]
+
+        # Shift left to reduce padding
+        min_left_pad = self.left_padding.min().item()
+        if min_left_pad > 0:
+            self.keys = self.keys[..., min_left_pad:, :]
+            self.values = self.values[..., min_left_pad:, :]
+            self._idx -= min_left_pad
+            self.left_padding -= min_left_pad
+
+    def extend(self, other):
+        """
+        In-place extend this cache with the other cache.
+        """
+        max_idx = max(self._idx, other._idx)
+        max_size = max(self.keys.shape[2], other.keys.shape[2])
+
+        # Pad the keys and values so they are right-justified
+        # with the index and the same size
+        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
+
+
+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
+
+    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, :]
+
+            # 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):
+        # 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)
+
+    @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)

mlx_lm/models/cohere.py

@@ -1,7 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.
 
 from dataclasses import dataclass
-from typing import Any, Optional, Tuple
+from typing import Any, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -155,17 +155,15 @@ class CohereModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
         h = self.embed_tokens(inputs)
 
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-
         if cache is None:
             cache = [None] * len(self.layers)
 
+        mask = create_attention_mask(h, cache[0])
+
         for layer, c in zip(self.layers, cache):
             h = layer(h, mask, c)
 
@@ -182,10 +180,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, cache)
         out = self.model.embed_tokens.as_linear(out)
         out = out * self.model.args.logit_scale
         return out

mlx_lm/models/cohere2.py

@@ -83,11 +83,6 @@ class Attention(nn.Module):
         if cache is not None:
             keys, values = cache.update_and_fetch(keys, values)
 
-        if self.use_sliding_window and isinstance(mask, mx.array):
-            key_len = keys.shape[-2]
-            if mask.shape[-1] != key_len:
-                mask = mask[..., -key_len:]
-
         # TODO: maybe remove cast once fused mask is supported since attention
         # may be in higher precision
         sdpa_type = mx.float32 if queries.dtype == mx.float16 else queries.dtype
@@ -148,6 +143,7 @@ class CohereModel(nn.Module):
         self.vocab_size = args.vocab_size
         self.num_hidden_layers = args.num_hidden_layers
         assert self.vocab_size > 0
+        self.window_size = args.sliding_window
         self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
         self.layers = [
             TransformerBlock(args=args, layer_idx=i)
@@ -160,7 +156,6 @@ class CohereModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
         h = self.embed_tokens(inputs)
@@ -168,10 +163,9 @@ class CohereModel(nn.Module):
         if cache is None:
             cache = [None] * len(self.layers)
 
-        if mask is None:
-            j = self.args.sliding_window_pattern
-            full_mask = create_attention_mask(h, cache[j - 1 : j])
-            sliding_window_mask = create_attention_mask(h, cache)
+        j = self.args.sliding_window_pattern
+        full_mask = create_attention_mask(h, cache[j - 1])
+        swa_mask = create_attention_mask(h, cache[0], window_size=self.window_size)
 
         for i, (layer, c) in enumerate(zip(self.layers, cache)):
             is_global = (
@@ -179,13 +173,9 @@ class CohereModel(nn.Module):
                 == self.args.sliding_window_pattern - 1
             )
 
-            local_mask = mask
-            if mask is None and is_global:
-                local_mask = full_mask
-            elif mask is None:
-                local_mask = sliding_window_mask
+            mask = full_mask if is_global else swa_mask
 
-            h = layer(h, local_mask, c)
+            h = layer(h, mask, c)
 
         return self.norm(h)
 
@@ -200,10 +190,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, cache)
         out = self.model.embed_tokens.as_linear(out)
         out = out * self.model.args.logit_scale
         return out

mlx_lm/models/dbrx.py

@@ -1,7 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.
 
 from dataclasses import dataclass
-from typing import Any, Optional, Tuple
+from typing import Any, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -105,10 +105,9 @@ class MLP(nn.Module):
         self.v1 = nn.Linear(d_model, ffn_dim, bias=False)
         self.w1 = nn.Linear(d_model, ffn_dim, bias=False)
         self.w2 = nn.Linear(ffn_dim, d_model, bias=False)
-        self.act_fn = nn.silu
 
     def __call__(self, x: mx.array) -> mx.array:
-        current_hidden_states = self.act_fn(self.w1(x)) * self.v1(x)
+        current_hidden_states = nn.silu(self.w1(x)) * self.v1(x)
         current_hidden_states = self.w2(current_hidden_states)
         return current_hidden_states
 
@@ -197,17 +196,15 @@ class DBRX(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
         h = self.wte(inputs)
 
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-
         if cache is None:
             cache = [None] * len(self.blocks)
 
+        mask = create_attention_mask(h, cache[0])
+
         for layer, c in zip(self.blocks, cache):
             h = layer(h, mask, c)
 
@@ -225,10 +222,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        out = self.transformer(inputs, mask, cache)
+        out = self.transformer(inputs, cache)
         return self.lm_head(out)
 
     @property

mlx_lm/models/deepseek.py

@@ -118,10 +118,9 @@ class DeepseekMLP(nn.Module):
         self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
         self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
         self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
-        self.act_fn = nn.silu
 
     def __call__(self, x: mx.array) -> mx.array:
-        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
 
 
 class MoEGate(nn.Module):
@@ -211,15 +210,14 @@ class DeepseekModel(nn.Module):
         self,
         x: mx.array,
         cache: Optional[Any] = None,
-        mask: Optional[mx.array] = None,
     ) -> mx.array:
         h = self.embed_tokens(x)
-        if mask is None:
-            mask = create_attention_mask(h, cache)
 
         if cache is None:
             cache = [None] * len(self.layers)
 
+        mask = create_attention_mask(h, cache[0])
+
         for layer, c in zip(self.layers, cache):
             h = layer(h, mask, c)
 
@@ -238,9 +236,8 @@ class Model(nn.Module):
         self,
         inputs: mx.array,
         cache: Optional[Any] = None,
-        mask: Optional[mx.array] = None,
     ):
-        out = self.model(inputs, cache, mask)
+        out = self.model(inputs, cache)
         return self.lm_head(out)
 
     def sanitize(self, weights):

mlx_lm/models/deepseek_v2.py

@@ -2,7 +2,7 @@
 
 import math
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -394,17 +394,15 @@ class DeepseekV2Model(nn.Module):
         self,
         x: mx.array,
         cache: Optional[Any] = None,
-        mask: Optional[mx.array] = None,
     ) -> mx.array:
         h = self.embed_tokens(x)
 
         pipeline_rank = self.pipeline_rank
         pipeline_size = self.pipeline_size
-        if mask is None:
-            mask = create_attention_mask(h, cache)
 
         if cache is None:
             cache = [None] * self.num_layers
+        mask = create_attention_mask(h, cache[0])
 
         # Receive from the previous process in the pipeline
         if pipeline_rank < pipeline_size - 1:
@@ -416,6 +414,8 @@ class DeepseekV2Model(nn.Module):
         # 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]]
@@ -435,9 +435,8 @@ class Model(nn.Module):
         self,
         inputs: mx.array,
         cache: Optional[Any] = None,
-        mask: Optional[mx.array] = None,
     ):
-        out = self.model(inputs, cache, mask)
+        out = self.model(inputs, cache)
         return self.lm_head(out)
 
     def sanitize(self, weights):

mlx_lm/models/deepseek_v3.py

@@ -3,7 +3,7 @@
 import math
 from dataclasses import dataclass
 from functools import partial
-from typing import Any, Dict, Optional, Tuple
+from typing import Any, Dict, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -33,9 +33,9 @@ class ModelArgs(BaseModelArgs):
     topk_method: str = "noaux_tc"
     scoring_func: str = "sigmoid"
     norm_topk_prob: bool = True
-    n_group: Optional[int] = None
-    topk_group: Optional[int] = None
-    num_experts_per_tok: Optional[int] = None
+    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
@@ -124,12 +124,6 @@ class DeepseekV3YarnRotaryEmbedding(nn.Module):
         )
 
 
-# A clipped silu to prevent fp16 from overflowing
-@partial(mx.compile, shapeless=True)
-def clipped_silu(x):
-    return mx.clip(x * mx.sigmoid(x), a_min=-100, a_max=100)
-
-
 class DeepseekV3Attention(nn.Module):
     def __init__(self, config: ModelArgs):
         super().__init__()
@@ -287,16 +281,18 @@ def group_expert_select(
     norm_topk_prob,
 ):
 
-    k = top_k
     scores = mx.sigmoid(gates.astype(mx.float32))
     orig_scores = scores
     scores = scores + e_score_correction_bias
-    scores = mx.unflatten(scores, axis=-1, shape=(n_group, -1))
-    group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True)
-    k = n_group - topk_group
-    group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :]
-    scores = mx.put_along_axis(scores, group_idx, mx.array(0.0), axis=-2)
-    scores = mx.flatten(scores, -2, -1)
+    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]
@@ -344,7 +340,6 @@ class DeepseekV3MoE(nn.Module):
             config.hidden_size,
             config.moe_intermediate_size,
             config.n_routed_experts,
-            activation=clipped_silu,
         )
 
         self.gate = MoEGate(config)
@@ -430,17 +425,15 @@ class DeepseekV3Model(nn.Module):
         self,
         x: mx.array,
         cache: Optional[Any] = None,
-        mask: Optional[mx.array] = None,
     ) -> mx.array:
         h = self.embed_tokens(x)
 
         pipeline_rank = self.pipeline_rank
         pipeline_size = self.pipeline_size
-        if mask is None:
-            mask = create_attention_mask(h, cache)
 
         if cache is None:
             cache = [None] * self.num_layers
+        mask = create_attention_mask(h, cache[0])
 
         # Receive from the previous process in the pipeline
 
@@ -453,6 +446,8 @@ class DeepseekV3Model(nn.Module):
         # 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]]
@@ -472,9 +467,8 @@ class Model(nn.Module):
         self,
         inputs: mx.array,
         cache: Optional[Any] = None,
-        mask: Optional[mx.array] = None,
     ):
-        out = self.model(inputs, cache, mask)
+        out = self.model(inputs, cache)
         return self.lm_head(out)
 
     def sanitize(self, weights):
@@ -529,6 +523,7 @@ class Model(nn.Module):
     def layers(self):
         return self.model.layers[self.model.start_idx : self.model.end_idx]
 
+    @property
     def cast_predicate(self):
         def predicate(k):
             return "e_score_correction_bias" not in k

mlx_lm/models/dots1.py

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

mlx_lm/models/ernie4_5.py

@@ -0,0 +1,164 @@
+# Copyright © 2023-2024 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    hidden_size: int
+    intermediate_size: int
+    model_type: str
+    max_position_embeddings: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    head_dim: Optional[int]
+    num_hidden_layers: int
+    rms_norm_eps: float
+    vocab_size: int
+    rope_theta: float
+    use_bias: bool
+    tie_word_embeddings: bool
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+
+        self.head_dim = head_dim = args.head_dim or dim // n_heads
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=args.use_bias)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.use_bias)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.use_bias)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=args.use_bias)
+
+        self.rope = initialize_rope(
+            head_dim,
+            base=args.rope_theta,
+            traditional=True,
+            max_position_embeddings=args.max_position_embeddings,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, dim, hidden_dim, use_bias=False):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=use_bias)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=use_bias)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=use_bias)
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.self_attn = Attention(args)
+        self.mlp = MLP(args.hidden_size, args.intermediate_size, args.use_bias)
+
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        return h + r
+
+
+class Ernie45Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [DecoderLayer(args) for _ in range(args.num_hidden_layers)]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+    ):
+        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 = Ernie45Model(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        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
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/ernie4_5_moe.py

@@ -0,0 +1,288 @@
+# Copyright © 2023-2024 Apple Inc.
+
+from dataclasses import dataclass, field
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    hidden_size: int
+    intermediate_size: int
+    model_type: str
+    max_position_embeddings: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    num_hidden_layers: int
+    rms_norm_eps: float
+    vocab_size: int
+    rope_theta: float
+    use_bias: bool
+    tie_word_embeddings: bool
+    moe_num_experts: int
+    moe_layer_start_index: int = 0
+    moe_intermediate_size: int = 0
+    moe_capacity: list[int] = field(default_factory=list)
+    moe_k: int = 1
+    moe_layer_interval: int = 1
+    moe_use_aux_free: bool = False
+    moe_num_shared_experts: int = 0
+    moe_layer_end_index: Optional[int] = None
+    head_dim: Optional[int] = None
+    moe_gate_act: str = "softmax"
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+
+        self.head_dim = head_dim = args.head_dim or dim // n_heads
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=args.use_bias)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.use_bias)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.use_bias)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=args.use_bias)
+
+        self.rope = initialize_rope(
+            head_dim,
+            base=args.rope_theta,
+            traditional=True,
+            max_position_embeddings=args.max_position_embeddings,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class Ernie4_5_MLP(nn.Module):
+    def __init__(self, dim, hidden_dim, use_bias=False):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=use_bias)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=use_bias)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=use_bias)
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class Ernie4_5_MoeMLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.k = args.moe_k
+        self.moe_intermediate_size = (
+            args.moe_intermediate_size
+            if args.moe_intermediate_size
+            else args.intermediate_size
+        )
+
+        self.gate = nn.Linear(args.hidden_size, args.moe_num_experts, bias=False)
+
+        self.switch_mlp = SwitchGLU(
+            args.hidden_size,
+            self.moe_intermediate_size,
+            args.moe_num_experts,
+            bias=args.use_bias,
+        )
+
+        if getattr(args, "moe_num_shared_experts", 0) > 0:
+            shared_intermediate_size = (
+                args.moe_intermediate_size * args.moe_num_shared_experts
+                if getattr(args, "moe_intermediate_size", None)
+                else args.intermediate_size * args.moe_num_shared_experts
+            )
+            self.shared_experts = Ernie4_5_MLP(
+                args.hidden_size, shared_intermediate_size, args.use_bias
+            )
+        else:
+            self.shared_experts = None
+
+        if args.moe_gate_act == "softmax":
+            self.gate_act = nn.Softmax()
+        elif args.moe_gate_act == "sigmoid":
+            self.gate_act = nn.Sigmoid()
+        else:
+            raise ValueError(f"{args.moe_gate_act} is not supported.")
+
+    def __call__(self, x: mx.array) -> mx.array:
+        gates = self.gate(x)
+        gates = self.gate_act(gates.astype(mx.float32))
+
+        k = self.k
+        inds = mx.stop_gradient(mx.argpartition(-gates, kth=k - 1, axis=-1)[..., :k])
+        scores = mx.take_along_axis(gates, inds, axis=-1)
+
+        scores = scores / mx.maximum(scores.sum(axis=-1, keepdims=True), 1e-12)
+
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
+
+        if self.shared_experts is not None:
+            y = y + self.shared_experts(x)
+
+        return y
+
+
+class Ernie4_5_DecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.self_attn = Attention(args)
+
+        moe_layer_start_index = (
+            min(args.moe_layer_start_index)
+            if isinstance(args.moe_layer_start_index, (tuple, list))
+            else args.moe_layer_start_index
+        )
+
+        if args.moe_layer_end_index is None:
+            moe_layer_end_index = args.num_hidden_layers - 1
+        else:
+            moe_layer_end_index = (
+                max(args.moe_layer_end_index)
+                if isinstance(args.moe_layer_end_index, (tuple, list))
+                else args.moe_layer_end_index
+            )
+
+        if (
+            ((layer_idx + 1) % args.moe_layer_interval == 0)
+            and layer_idx >= moe_layer_start_index
+            and layer_idx <= moe_layer_end_index
+        ):
+            self.mlp = Ernie4_5_MoeMLP(args)
+        else:
+            self.mlp = Ernie4_5_MLP(
+                args.hidden_size, args.intermediate_size, args.use_bias
+            )
+
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        return h + r
+
+
+class Ernie45Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            Ernie4_5_DecoderLayer(args, i) for i in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+    ):
+        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 = Ernie45Model(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        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
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def sanitize(self, weights):
+        remove_patterns = [
+            "mtp_block.",
+            "mtp_linear_proj.",
+            "mtp_hidden_norm.",
+            "mtp_emb_norm.",
+            "e_score_correction_bias",
+        ]
+
+        weights = {
+            key: value
+            for key, value in weights.items()
+            if not any(pattern in key for pattern in remove_patterns)
+        }
+
+        # Stack experts
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}"
+            for m in ["gate_proj", "down_proj", "up_proj"]:
+                if f"{prefix}.mlp.experts.0.{m}.weight" in weights:
+                    to_join = [
+                        weights.pop(f"{prefix}.mlp.experts.{e}.{m}.weight")
+                        for e in range(self.args.moe_num_experts)
+                    ]
+                    weights[f"{prefix}.mlp.switch_mlp.{m}.weight"] = mx.stack(to_join)
+
+        return weights

mlx_lm/models/exaone.py

@@ -123,16 +123,15 @@ class ExaoneModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
         h = self.wte(inputs)
-        if mask is None:
-            mask = create_attention_mask(h, cache)
 
         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, cache=c)
 
@@ -151,10 +150,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        out = self.transformer(inputs, mask, cache)
+        out = self.transformer(inputs, cache)
         if self.args.tie_word_embeddings:
             out = self.transformer.wte.as_linear(out)
         else:

mlx_lm/models/exaone4.py

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

mlx_lm/models/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

@@ -0,0 +1,251 @@
+from functools import partial
+from typing import Optional, Tuple
+
+import mlx.core as mx
+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)
+    )
+
+
+def _make_gated_delta_kernel(has_mask=False):
+    if not mx.metal.is_available():
+        return None
+    mask_source = "mask[b_idx * T + t]" if has_mask else "true"
+    source = f"""
+        auto n = thread_position_in_grid.z;
+        auto b_idx = n / Hv;
+        auto hv_idx = n % Hv;
+        auto hk_idx = hv_idx / (Hv / Hk);
+        constexpr int n_per_t = Dk / 32;
+
+        // q, k: [B, T, Hk, Dk]
+        auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk;
+        auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk;
+
+        // v, y: [B, T, Hv, Dv]
+        auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv;
+        y += b_idx * T * Hv * Dv + hv_idx * Dv;
+
+        auto dk_idx = thread_position_in_threadgroup.x;
+        auto dv_idx = thread_position_in_grid.y;
+
+        // state_in, state_out: [B, Hv, Dv, Dk]
+        auto i_state = state_in + (n * Dv + dv_idx) * Dk;
+        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]);
+        }}
+
+        // beta, g: [B, T, Hv]
+        auto g_ = g + b_idx * T * Hv;
+        auto beta_ = beta + b_idx * T * Hv;
+
+        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_[hv_idx];
+              kv_mem += state[i] * k_[s_idx];
+            }}
+            kv_mem = simd_sum(kv_mem);
+
+            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);
+            }}
+          }}
+          // Increment data pointers to next time step
+          q_ += Hk * Dk;
+          k_ += Hk * Dk;
+          v_ += Hv * Dv;
+          y += Hv * Dv;
+          g_ += Hv;
+          beta_ += Hv;
+        }}
+        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")
+    return mx.fast.metal_kernel(
+        name="gated_delta_step" + "_mask" if has_mask else "",
+        input_names=inputs,
+        output_names=["y", "state_out"],
+        source=source,
+    )
+
+
+_gated_delta_kernel = _make_gated_delta_kernel()
+_gated_delta_kernel_masked = _make_gated_delta_kernel(True)
+
+
+@mx.compile
+def _gated_delta_step_ops(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    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.
+
+    Shapes:
+      - q, k: [B, H, Dk]
+      - v: [B, H, Dv]
+      - g, beta: [B, H]
+      - state: [B, H, Dv, Dk]
+    Returns:
+      - y: [B, H, Dv]
+      - new_state: [B, H, Dv, Dk]
+    """
+
+    # Decay
+    old_state = state
+    state = state * g[..., None, None]
+    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:
+        state = mx.where(mask, state, old_state)
+    return y, state
+
+
+def gated_delta_kernel(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    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
+    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),
+            ("Dv", Dv),
+            ("Hk", Hk),
+            ("Hv", Hv),
+        ],
+        grid=(32, Dv, B * Hv),
+        threadgroup=(32, 4, 1),
+        output_shapes=[(B, T, Hv, Dv), state.shape],
+        output_dtypes=[input_type, input_type],
+    )
+
+
+def gated_delta_ops(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    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).
+
+    Shapes:
+      - q, k: [B, T, Hk, Dk]
+      - v: [B, T, Hv, Dv]
+      - g, beta: [B, T, Hv]
+      - state: [B, Hv, Dk, Dv]
+    Returns:
+      - y: [B, T, Hv, Dv]
+      - state: [B, Hv, Dk, Dv]
+    """
+    B, T, Hk, Dk = q.shape
+    Hv, Dv = v.shape[-2:]
+    if state is None:
+        state = mx.zeros((B, Hv, Dv, Dk), dtype=q.dtype)
+
+    if (repeat_factor := Hv // Hk) > 1:
+        q = mx.repeat(q, repeat_factor, -2)
+        k = mx.repeat(k, repeat_factor, -2)
+
+    ys = []
+    for t in range(T):
+        if mask is not None:
+            y, state = _gated_delta_step_ops(
+                q[:, t],
+                k[:, t],
+                v[:, t],
+                g[:, t],
+                beta[:, t],
+                state,
+                mask[:, t],
+            )
+        else:
+            y, state = _gated_delta_step_ops(
+                q[:, t],
+                k[:, t],
+                v[:, t],
+                g[:, t],
+                beta[:, t],
+                state,
+            )
+        ys.append(y)
+    y = mx.stack(ys, axis=1)
+    return y, state
+
+
+def gated_delta_update(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    a: mx.array,
+    b: mx.array,
+    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)
+    g = compute_g(A_log, a, dt_bias)
+    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)
+
+    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)
+    else:
+        return gated_delta_kernel(q, k, v, g, beta, state, mask)

mlx_lm/models/gemma.py

@@ -1,7 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.
 
 from dataclasses import dataclass
-from typing import Any, Optional, Tuple
+from typing import Any, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -138,18 +138,16 @@ class GemmaModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
         h = self.embed_tokens(inputs)
         h = h * (self.args.hidden_size**0.5)
 
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-
         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)
 
@@ -166,10 +164,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, cache)
         out = self.model.embed_tokens.as_linear(out)
         return out
 

mlx_lm/models/gemma2.py

@@ -1,7 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.
 
 from dataclasses import dataclass
-from typing import Any, Optional, Tuple
+from typing import Any, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -165,18 +165,16 @@ class GemmaModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
         h = self.embed_tokens(inputs)
         h = h * (self.args.hidden_size**0.5)
 
-        if mask is None:
-            mask = create_attention_mask(h, cache, return_array=True)
-
         if cache is None:
             cache = [None] * len(self.layers)
 
+        mask = create_attention_mask(h, cache[0], return_array=True)
+
         for layer, c in zip(self.layers, cache):
             h = layer(h, mask, c)
 
@@ -194,10 +192,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, cache)
         out = self.model.embed_tokens.as_linear(out)
         out = mx.tanh(out / self.final_logit_softcapping)
         out = out * self.final_logit_softcapping

mlx_lm/models/gemma3.py

@@ -40,9 +40,11 @@ class Model(nn.Module):
         self,
         inputs: mx.array,
         cache=None,
-        mask: Optional[mx.array] = None,
+        input_embeddings: Optional[mx.array] = None,
     ):
-        return self.language_model(inputs, cache=cache, mask=mask)
+        return self.language_model(
+            inputs, cache=cache, input_embeddings=input_embeddings
+        )
 
     def sanitize(self, weights):
         weights = tree_unflatten(list(weights.items()))

mlx_lm/models/gemma3_text.py

@@ -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
         )
@@ -160,6 +158,8 @@ class Gemma3Model(nn.Module):
     def __init__(self, args: ModelArgs):
         super().__init__()
         self.args = args
+        self.window_size = args.sliding_window
+        self.sliding_window_pattern = args.sliding_window_pattern
         self.vocab_size = args.vocab_size
         self.num_hidden_layers = args.num_hidden_layers
         assert self.vocab_size > 0
@@ -173,34 +173,31 @@ class Gemma3Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
+        input_embeddings: Optional[mx.array] = None,
     ):
-
-        h = self.embed_tokens(inputs)
+        if input_embeddings is not None:
+            h = input_embeddings
+        else:
+            h = self.embed_tokens(inputs)
         h *= mx.array(self.args.hidden_size**0.5, mx.bfloat16).astype(h.dtype)
 
         if cache is None:
             cache = [None] * len(self.layers)
 
-        if mask is None:
-            j = self.args.sliding_window_pattern
-            full_mask = create_attention_mask(h, cache[j - 1 : j])
-            sliding_window_mask = create_attention_mask(h, cache)
+        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,
+        )
         for i, (layer, c) in enumerate(zip(self.layers, cache)):
             is_global = (
-                i % self.args.sliding_window_pattern
-                == self.args.sliding_window_pattern - 1
+                i % self.sliding_window_pattern == self.sliding_window_pattern - 1
             )
-
-            local_mask = mask
-            if mask is None and is_global:
-                local_mask = full_mask
-            elif mask is None:
-                local_mask = sliding_window_mask
-
-            h = layer(h, local_mask, c)
+            mask = global_mask if is_global else sliding_window_mask
+            h = layer(h, mask, c)
 
         return self.norm(h)
 
@@ -212,21 +209,25 @@ class Model(nn.Module):
         self.model_type = args.model_type
         self.model = Gemma3Model(args)
         self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+        self.tie_word_embeddings = False
 
     def __call__(
         self,
         inputs: mx.array,
         cache=None,
-        mask: Optional[mx.array] = None,
+        input_embeddings: Optional[mx.array] = None,
     ):
-        out = self.model(inputs, mask, cache)
-        out = self.lm_head(out)
+        out = self.model(inputs, cache, input_embeddings)
+        if self.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
         return out
 
     def sanitize(self, weights):
-        weights = dict(weights)
         if "lm_head.weight" not in weights:
-            weights["lm_head.weight"] = weights["model.embed_tokens.weight"]
+            self.tie_word_embeddings = True
+            self.pop("lm_head")
         return weights
 
     @property
@@ -242,7 +243,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/gemma3n.py

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

mlx_lm/models/glm.py

@@ -0,0 +1,187 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    rms_norm_eps: float
+    vocab_size: int
+    head_dim: int
+    num_key_value_heads: int
+    max_position_embeddings: Optional[int] = None
+    attention_bias: bool = False
+    rope_theta: float = 10000
+    tie_word_embeddings: bool = True
+
+
+class GLMAttention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.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 = args.head_dim or args.hidden_size // self.num_attention_heads
+        self.scale = self.head_dim**-0.5
+
+        self.q_proj = nn.Linear(
+            self.hidden_size,
+            self.num_attention_heads * self.head_dim,
+            bias=args.attention_bias,
+        )
+        self.k_proj = nn.Linear(
+            self.hidden_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=args.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            self.hidden_size,
+            self.num_key_value_heads * self.head_dim,
+            bias=args.attention_bias,
+        )
+        self.o_proj = nn.Linear(
+            self.num_attention_heads * self.head_dim, self.hidden_size, bias=False
+        )
+
+        self.rope = nn.RoPE(dims=self.head_dim, traditional=True, base=args.rope_theta)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        queries = queries.reshape(B, L, self.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 GLMMLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.gate_up_proj = nn.Linear(
+            args.hidden_size, 2 * args.intermediate_size, bias=False
+        )
+        self.down_proj = nn.Linear(args.intermediate_size, args.hidden_size, bias=False)
+
+    def __call__(self, x) -> mx.array:
+        x = self.gate_up_proj(x)
+        gate, x = mx.split(x, 2, axis=-1)
+        return self.down_proj(nn.silu(gate) * x)
+
+
+class GLMBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.self_attn = GLMAttention(args)
+        self.mlp = GLMMLP(args)
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        out = h + r
+        return out
+
+
+class GLMModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [GLMBlock(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,
+        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 = GLMModel(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 sanitize(self, weights):
+        weights = {
+            k: v for k, v in weights.items() if "self_attn.rotary_emb.inv_freq" not in k
+        }
+        if self.args.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
+        return weights
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/glm4.py

@@ -144,17 +144,15 @@ class Glm4Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: Optional[mx.array] = None,
         cache: Optional[Any] = None,
     ):
         h = self.embed_tokens(inputs)
 
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-
         if cache is None:
             cache = [None] * len(self.layers)
 
+        mask = create_attention_mask(h, cache[0])
+
         for layer, c in zip(self.layers, cache):
             h = layer(h, mask, cache=c)
 
@@ -172,10 +170,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: Optional[mx.array] = None,
         cache: Optional[Any] = None,
     ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, cache)
         return self.lm_head(out)
 
     @property

mlx_lm/models/glm4_moe.py

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

mlx_lm/models/gpt2.py

@@ -5,7 +5,6 @@ from typing import Any, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
-import numpy as np
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
@@ -126,25 +125,23 @@ class GPT2Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
         _, L = inputs.shape
 
         hidden_states = self.wte(inputs)
 
+        if cache is None:
+            cache = [None] * len(self.h)
+
         offset = 0
-        if cache is not None and len(cache) > 0 and cache[0] is not None:
+        if cache[0] is not None:
             offset = cache[0].offset
 
         position_ids = mx.arange(offset, offset + L)
         hidden_states += self.wpe(position_ids)
 
-        if mask is None:
-            mask = create_attention_mask(hidden_states, cache)
-
-        if cache is None:
-            cache = [None] * len(self.h)
+        mask = create_attention_mask(hidden_states, cache[0])
 
         for layer, c in zip(self.h, cache):
             hidden_states = layer(hidden_states, mask, cache=c)
@@ -162,10 +159,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, cache)
         out = self.model.wte.as_linear(out)
         return out
 

mlx_lm/models/gpt_bigcode.py

@@ -1,7 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.
 
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -137,23 +137,20 @@ class GPTBigCodeModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
         B, L = inputs.shape
 
         hidden_states = self.wte(inputs)
 
-        mask = None
-        if mask is not None and hidden_states.shape[1] > 1:
-            mask = create_attention_mask(hidden_states, cache)
-
         if cache is None:
             cache = [None] * len(self.h)
             position_ids = mx.array(np.arange(L))
         else:
             position_ids = mx.array(np.arange(cache[0].offset, cache[0].offset + L))
 
+        mask = create_attention_mask(hidden_states, cache[0])
+
         hidden_states += self.wpe(position_ids)
 
         for layer, c in zip(self.h, cache):
@@ -174,10 +171,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        out = self.transformer(inputs, mask, cache)
+        out = self.transformer(inputs, cache)
         if self.args.tie_word_embeddings:
             out = self.transformer.wte.as_linear(out)
         else:

mlx_lm/models/gpt_neox.py

@@ -1,11 +1,10 @@
 # Copyright © 2023-2024 Apple Inc.
 
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
-import numpy as np
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
@@ -146,19 +145,17 @@ class GPTNeoXModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
         _, L = inputs.shape
 
         hidden_states = self.embed_in(inputs)
 
-        if mask is None:
-            mask = create_attention_mask(hidden_states, cache)
-
         if cache is None:
             cache = [None] * len(self.h)
 
+        mask = create_attention_mask(hidden_states, cache[0])
+
         for layer, c in zip(self.h, cache):
             hidden_states = layer(hidden_states, mask, cache=c)
 
@@ -178,10 +175,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, cache)
         return out
 
     def sanitize(self, weights):

mlx_lm/models/gpt_oss.py

@@ -0,0 +1,291 @@
+# 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
+from .cache import KVCache, RotatingKVCache
+from .rope_utils import initialize_rope
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str = "gpt_oss"
+    num_hidden_layers: int = 36
+    num_local_experts: int = 128
+    num_experts_per_tok: int = 4
+    vocab_size: int = 201088
+    rms_norm_eps: float = 1e-05
+    hidden_size: int = 2880
+    intermediate_size: int = 2880
+    head_dim: int = 64
+    num_attention_heads: int = 64
+    num_key_value_heads: int = 8
+    sliding_window: int = 128
+    rope_theta: int = 150000
+    rope_scaling: Any = None
+    layer_types: list = None
+
+
+# These operators emulate particular methods in torch that don't exist in MLX natively
+def mlx_topk(a, k, axis=-1):
+    """MLX equivalent of torch.topk"""
+    partitioned_indices = mx.argpartition(a, kth=-k, axis=axis)
+    # Extract only the top k indices (last k elements after partition)
+    top_k_indices = partitioned_indices[..., -k:]
+    # Get the corresponding values
+    top_k_values = mx.take_along_axis(a, top_k_indices, axis=axis)
+    return top_k_values, top_k_indices
+
+
+@partial(mx.compile, shapeless=True)
+def swiglu(x_linear, x_glu, alpha: float = 1.702, limit: float = 7.0):
+    # Clamp the input values
+    x_glu = mx.clip(x_glu, a_min=None, a_max=limit)
+    x_linear = mx.clip(x_linear, a_min=-limit, a_max=limit)
+
+    glu_scaled = alpha * x_glu
+    sig = mx.sigmoid(glu_scaled)
+
+    out_glu = x_glu * sig
+    # Note we add an extra bias of 1 to the linear layer
+    return out_glu * (x_linear + 1)
+
+
+class SwiGLU(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def __call__(self, x, gate):
+        return swiglu(x, gate)
+
+
+class AttentionBlock(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+
+        self.head_dim = config.head_dim
+        self.num_attention_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = (
+            config.num_attention_heads // config.num_key_value_heads
+        )
+
+        self.sinks = mx.zeros((config.num_attention_heads,))
+
+        self.q_proj = nn.Linear(
+            config.hidden_size, config.num_attention_heads * self.head_dim, bias=True
+        )
+        self.k_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=True
+        )
+        self.v_proj = nn.Linear(
+            config.hidden_size, config.num_key_value_heads * self.head_dim, bias=True
+        )
+
+        self.o_proj = nn.Linear(
+            self.head_dim * config.num_attention_heads, config.hidden_size, bias=True
+        )
+
+        self.sm_scale = 1 / math.sqrt(config.head_dim)
+
+        self.rope = initialize_rope(
+            self.head_dim,
+            config.rope_theta,
+            traditional=False,
+            scaling_config=config.rope_scaling,
+        )
+
+    def __call__(self, x: mx.array, mask: mx.array, cache=None) -> mx.array:
+        B, L, _ = x.shape
+        D = self.head_dim
+        Hk = self.num_key_value_heads
+
+        q = self.q_proj(x).reshape(B, L, -1, D).swapaxes(1, 2)
+        k = self.k_proj(x).reshape(B, L, -1, D).swapaxes(1, 2)
+        v = self.v_proj(x).reshape(B, L, -1, D).swapaxes(1, 2)
+
+        if cache is not None:
+            q = self.rope(q, offset=cache.offset)
+            k = self.rope(k, offset=cache.offset)
+            k, v = cache.update_and_fetch(k, v)
+        else:
+            q = self.rope(q)
+            k = self.rope(k)
+
+        v_hat = scaled_dot_product_attention(
+            q, k, v, cache, self.sm_scale, mask=mask, sinks=self.sinks
+        )
+
+        return self.o_proj(v_hat.swapaxes(1, 2).reshape(B, L, -1))
+
+
+class MLPBlock(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+
+        self.hidden_size = config.hidden_size
+        self.num_local_experts = config.num_local_experts
+        self.num_experts_per_tok = config.num_experts_per_tok
+
+        self.experts = SwitchGLU(
+            input_dims=config.hidden_size,
+            hidden_dims=config.intermediate_size,
+            num_experts=config.num_local_experts,
+            activation=SwiGLU(),
+            bias=True,
+        )
+        self.router = nn.Linear(config.hidden_size, config.num_local_experts, bias=True)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        g = self.router(x)
+        experts, indices = mlx_topk(g, k=self.num_experts_per_tok, axis=-1)
+        expert_weights = mx.softmax(experts, axis=-1, precise=True)
+
+        # Experts block
+        x = self.experts(x, indices)
+
+        x = x * mx.expand_dims(expert_weights, axis=-1)
+        return x.sum(axis=-2)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.self_attn = AttentionBlock(config)
+        self.mlp = MLPBlock(config)
+        self.input_layernorm = nn.RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+
+    def __call__(self, x: mx.array, mask: mx.array, cache=None) -> mx.array:
+        residual = x
+        x = self.input_layernorm(x)
+        x = self.self_attn(x, mask, cache)
+        x = residual + x
+
+        residual = x
+        x = self.post_attention_layernorm(x)
+        x = self.mlp(x)
+        x = residual + x
+        return x
+
+
+class GptOssMoeModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.norm = nn.RMSNorm(args.hidden_size, args.rms_norm_eps)
+        self.layer_types = args.layer_types or [
+            "sliding_attention",
+            "full_attention",
+        ] * (args.num_hidden_layers // 2)
+        self.layers = [TransformerBlock(args) for _ in range(args.num_hidden_layers)]
+        self.window_size = args.sliding_window
+        self.swa_idx = self.layer_types.index("sliding_attention")
+        self.ga_idx = self.layer_types.index("full_attention")
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        if input_embeddings is not None:
+            x = input_embeddings
+        else:
+            x = self.embed_tokens(inputs)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        full_mask = create_attention_mask(x, cache[self.ga_idx])
+        swa_mask = create_attention_mask(
+            x, cache[self.swa_idx], window_size=self.window_size
+        )
+
+        for layer, c, layer_type in zip(self.layers, cache, self.layer_types):
+            mask = full_mask if layer_type == "full_attention" else swa_mask
+            x = layer(x, mask, c)
+        x = self.norm(x)
+        return x
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = GptOssMoeModel(args)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(self, inputs: mx.array, cache=None):
+        return self.lm_head(self.model(inputs, cache))
+
+    def sanitize(self, weights):
+        if any("gate_proj.weight" in k for k in weights.keys()):
+            return weights  # already sanitized
+
+        new_weights = {}
+        for k, v in weights.items():
+            if "gate_up_proj" in k and "bias" not in k:
+                if "_blocks" in k:
+                    v = v.view(mx.uint32).flatten(-2)
+                    k = k.replace("_blocks", ".weight")
+                if "_scales" in k:
+                    k = k.replace("_scales", ".scales")
+                new_weights[k.replace("gate_up_proj", "gate_proj")] = mx.contiguous(
+                    v[..., ::2, :]
+                )
+                new_weights[k.replace("gate_up_proj", "up_proj")] = mx.contiguous(
+                    v[..., 1::2, :]
+                )
+            elif "down_proj" in k and "bias" not in k:
+                if "_blocks" in k:
+                    v = v.view(mx.uint32).flatten(-2)
+                    k = k.replace("_blocks", ".weight")
+                if "_scales" in k:
+                    k = k.replace("_scales", ".scales")
+                new_weights[k] = v
+            elif "gate_up_proj_bias" in k:
+                new_weights[k.replace("gate_up_proj_bias", "gate_proj.bias")] = (
+                    mx.contiguous(v[..., ::2])
+                )
+                new_weights[k.replace("gate_up_proj_bias", "up_proj.bias")] = (
+                    mx.contiguous(v[..., 1::2])
+                )
+            elif "down_proj_bias" in k:
+                new_weights[k.replace("down_proj_bias", "down_proj.bias")] = v
+            else:
+                new_weights[k] = v
+
+        return new_weights
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if path.endswith("router"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate
+
+    def make_cache(self):
+        caches = []
+        for lt in self.model.layer_types:
+            if lt == "full_attention":
+                caches.append(KVCache())
+            else:
+                caches.append(RotatingKVCache(max_size=self.args.sliding_window))
+        return caches

mlx_lm/models/granite.py

@@ -150,17 +150,15 @@ class GraniteModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
         h = self.embed_tokens(inputs) * self.embedding_multiplier
 
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-
         if cache is None:
             cache = [None] * len(self.layers)
 
+        mask = create_attention_mask(h, cache[0])
+
         for layer, c in zip(self.layers, cache):
             h = layer(h, mask, cache=c)
 
@@ -180,10 +178,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, cache)
         if self.args.tie_word_embeddings:
             out = self.model.embed_tokens.as_linear(out)
         else:

mlx_lm/models/granitemoe.py

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

mlx_lm/models/granitemoehybrid.py

@@ -0,0 +1,541 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, 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 .rope_utils import initialize_rope
+from .ssm import ssm_update
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    # Required fields (no defaults)
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    num_hidden_layers: int
+    max_position_embeddings: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    attention_bias: bool
+    embedding_multiplier: float
+    attention_multiplier: float
+    logits_scaling: float
+    residual_multiplier: float
+    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):
+        super().__init__()
+        self.eps = eps
+        self.weight = mx.ones(hidden_size)
+
+    def __call__(self, hidden_states: mx.array, gate: mx.array = None) -> mx.array:
+        if gate is not None:
+            hidden_states = hidden_states * nn.silu(gate)
+        return mx.fast.rms_norm(hidden_states, self.weight, self.eps)
+
+
+class GraniteMoeHybridMamba2Mixer(nn.Module):
+    def __init__(self, args: ModelArgs):
+        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_n_heads * args.mamba_d_head
+        self.n_groups = args.mamba_n_groups
+        self.head_dim = args.mamba_d_head
+        self.time_step_limit = args.time_step_limit
+        self.heads_per_group = self.num_heads // self.n_groups
+
+        self.conv_dim = self.intermediate_size + 2 * self.n_groups * self.ssm_state_size
+
+        self.conv1d = nn.Conv1d(
+            in_channels=self.conv_dim,
+            out_channels=self.conv_dim,
+            kernel_size=args.mamba_d_conv,
+            padding=0,
+            groups=self.conv_dim,
+            bias=args.mamba_conv_bias,
+        )
+
+        projection_size = self.intermediate_size + self.conv_dim + self.num_heads
+        self.in_proj = nn.Linear(
+            self.hidden_size, projection_size, bias=args.mamba_proj_bias
+        )
+
+        self.dt_bias = mx.ones(self.num_heads)
+        self.A_log = mx.log(mx.arange(1, self.num_heads + 1, dtype=mx.float32))
+        self.D = mx.ones(self.num_heads)
+
+        self.norm = GraniteMoeHybridRMSNormGated(
+            self.intermediate_size, eps=args.rms_norm_eps
+        )
+        self.out_proj = nn.Linear(
+            self.intermediate_size, self.hidden_size, bias=args.mamba_proj_bias
+        )
+
+    def _apply_conv(
+        self, conv_input: mx.array, cache: Optional[MambaCache] = None
+    ) -> mx.array:
+        if cache is 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,
+        hidden_states: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[MambaCache] = None,
+    ) -> mx.array:
+
+        projected = self.in_proj(hidden_states)
+
+        gate, conv_input, dt = mx.split(
+            projected,
+            [self.intermediate_size, self.intermediate_size + self.conv_dim],
+            axis=-1,
+        )
+
+        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
+        y = self.norm(y, gate)
+        return self.out_proj(y)
+
+
+class GraniteMoeHybridAttention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+
+        self.head_dim = head_dim = args.hidden_size // n_heads
+
+        self.scale = args.attention_multiplier
+        attention_bias = args.attention_bias
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=attention_bias)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=attention_bias)
+
+        # Check if RoPE should be used based on position_embedding_type
+        # If position_embedding_type is "nope", don't use RoPE
+        use_rope = args.position_embedding_type != "nope"
+        if use_rope:
+            self.rope = initialize_rope(
+                self.head_dim,
+                args.rope_theta,
+                False,
+                None,  # rope_scaling
+                args.max_position_embeddings,
+            )
+        else:
+            self.rope = None
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[KVCache] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        # Apply RoPE only if enabled
+        if 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)
+        return self.o_proj(output)
+
+
+class GraniteMoeHybridTopKGating(nn.Module):
+    def __init__(self, input_size: int, num_experts: int, top_k: int):
+        super().__init__()
+        self.num_experts = num_experts
+        self.input_size = input_size
+        self.top_k = top_k
+        self.layer = nn.Linear(input_size, num_experts, bias=False)
+
+    def __call__(self, hidden_states: mx.array):
+        logits = self.layer(hidden_states)
+        top_k_idx = mx.argpartition(logits, kth=-self.top_k, axis=-1)[
+            ..., -self.top_k :
+        ]
+        top_k_logits = mx.take_along_axis(logits, top_k_idx, axis=-1)
+        top_k_gates = mx.softmax(top_k_logits, precise=True, axis=-1)
+        return top_k_idx, top_k_gates
+
+
+class GraniteMoeHybridMoE(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        self.input_size = args.hidden_size
+        self.hidden_size = args.intermediate_size
+        self.switch_mlp = SwitchGLU(
+            self.input_size, self.hidden_size, args.num_local_experts
+        )
+        self.router = GraniteMoeHybridTopKGating(
+            input_size=self.input_size,
+            num_experts=args.num_local_experts,
+            top_k=args.num_experts_per_tok,
+        )
+
+    def __call__(self, x: mx.array) -> mx.array:
+        token_ids, gates = self.router(x)
+        y = self.switch_mlp(x, token_ids)
+        return (y * gates[..., None]).sum(axis=-2)
+
+
+class GraniteMoeHybridSharedMLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.input_linear = nn.Linear(
+            args.hidden_size, args.shared_intermediate_size * 2, bias=False
+        )
+        self.output_linear = nn.Linear(
+            args.shared_intermediate_size, args.hidden_size, bias=False
+        )
+
+    def __call__(self, x: mx.array) -> mx.array:
+        gate, up = mx.split(self.input_linear(x), 2, axis=-1)
+        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)
+
+        if layer_type == "mamba":
+            self.mamba = GraniteMoeHybridMamba2Mixer(args)
+        elif layer_type == "attention":
+            self.self_attn = GraniteMoeHybridAttention(args)
+        else:
+            raise ValueError(f"Unknown layer type: {layer_type}")
+
+        # 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
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        # First block: either Mamba or Attention
+        residual = x
+        hidden_states = self.input_layernorm(x)
+
+        if self.layer_type == "mamba":
+            hidden_states = self.mamba(hidden_states, mask=mask, cache=cache)
+        else:
+            hidden_states = self.self_attn(hidden_states, mask=mask, cache=cache)
+
+        hidden_states = residual + hidden_states * self.residual_multiplier
+
+        # Second block: MoE + shared_mlp OR dense MLP
+        residual = hidden_states
+        normed = self.post_attention_layernorm(hidden_states)
+
+        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
+
+        return hidden_states
+
+
+class GraniteMoeHybridModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            GraniteMoeHybridLayer(args, layer_type) for layer_type in args.layer_types
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.embedding_multiplier = args.embedding_multiplier
+
+        # 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,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        hidden_states = self.embed_tokens(inputs) * self.embedding_multiplier
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        # Create masks based on what layer types exist
+        attn_mask = None
+        mamba_mask = None
+
+        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)
+
+        return self.norm(hidden_states)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = GraniteMoeHybridModel(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+        self.logits_scaling = args.logits_scaling
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        out = self.model(inputs, cache=cache)
+
+        if self.args.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
+
+        return out / self.logits_scaling
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def make_cache(self):
+        caches = []
+        for layer in self.layers:
+            if layer.layer_type == "mamba":
+                caches.append(MambaCache())
+            elif layer.layer_type == "attention":
+                caches.append(KVCache())
+        return caches
+
+    def sanitize(self, weights):
+        # Handle conv1d weights
+        for k, v in weights.items():
+            if "conv1d.weight" in k and v.shape[-1] != 1:
+                weights[k] = v.moveaxis(2, 1)
+
+        # 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"
+
+                input_weight = weights.pop(f"{prefix}.input_linear.weight")
+                _, expert_hidden, _ = input_weight.shape
+
+                # Split into gate and up projections (each half of expert_hidden)
+                gate_proj = input_weight[:, : expert_hidden // 2, :]
+                up_proj = input_weight[:, expert_hidden // 2 :, :]
+                weights[f"{prefix}.switch_mlp.gate_proj.weight"] = gate_proj
+                weights[f"{prefix}.switch_mlp.up_proj.weight"] = up_proj
+
+                weights[f"{prefix}.switch_mlp.down_proj.weight"] = weights.pop(
+                    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 self.args.use_moe and path.endswith("router.layer"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate

mlx_lm/models/helium.py

@@ -1,7 +1,7 @@
 # Copyright © 2025 Apple Inc.
 
 from dataclasses import dataclass
-from typing import Any, Optional, Tuple
+from typing import Any, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -136,17 +136,15 @@ class HeliumModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ) -> mx.array:
         h = self.embed_tokens(inputs)
 
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-
         if cache is None:
             cache = [None] * len(self.layers)
 
+        mask = create_attention_mask(h, cache[0])
+
         for layer, c in zip(self.layers, cache):
             h = layer(h, mask, c)
 
@@ -170,10 +168,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ) -> mx.array:
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, cache)
         if self.args.tie_word_embeddings:
             out = self.model.embed_tokens.as_linear(out)
         else:

mlx_lm/models/hunyuan.py

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

mlx_lm/models/hunyuan_v1_dense.py

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

mlx_lm/models/internlm2.py

@@ -1,7 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.
 
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Dict, Optional, Union
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -193,17 +193,14 @@ class InternLM2Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
         h = self.tok_embeddings(inputs)
 
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-
         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, cache=c)
 
@@ -222,10 +219,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, cache)
         if self.args.tie_word_embeddings:
             out = self.model.tok_embeddings.as_linear(out)
         else:

mlx_lm/models/internlm3.py

@@ -1,7 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.
 
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Dict, Optional, Union
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -193,17 +193,14 @@ class InternLM2Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
         h = self.embed_tokens(inputs)
 
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-
         if cache is None:
             cache = [None] * len(self.layers)
 
+        mask = create_attention_mask(h, cache[0])
         for layer, c in zip(self.layers, cache):
             h = layer(h, mask, cache=c)
 
@@ -222,10 +219,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, cache)
         if self.args.tie_word_embeddings:
             out = self.model.embed_tokens.as_linear(out)
         else:

mlx_lm/models/kimi_vl.py

@@ -1,16 +1,13 @@
 # Copyright © 2024 Apple Inc.
 
-import math
 from dataclasses import dataclass
-from functools import partial
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Dict, Optional, Union
 
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .base import BaseModelArgs
 from .deepseek_v3 import DeepseekV3Model
-from .switch_layers import SwitchGLU
 
 
 @dataclass
@@ -33,9 +30,9 @@ class TextArgs(BaseModelArgs):
     topk_method: str = "noaux_tc"
     scoring_func: str = "sigmoid"
     norm_topk_prob: bool = True
-    n_group: Optional[int] = None
-    topk_group: Optional[int] = None
-    num_experts_per_tok: Optional[int] = None
+    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
@@ -65,9 +62,8 @@ class LanguageModel(nn.Module):
         self,
         inputs: mx.array,
         cache: Optional[Any] = None,
-        mask: Optional[mx.array] = None,
     ):
-        out = self.model(inputs, cache, mask)
+        out = self.model(inputs, cache)
         return self.lm_head(out)
 
 
@@ -82,9 +78,8 @@ class Model(nn.Module):
         self,
         inputs: mx.array,
         cache: Optional[Any] = None,
-        mask: Optional[mx.array] = None,
     ):
-        return self.language_model(inputs, cache, mask)
+        return self.language_model(inputs, cache)
 
     def sanitize(self, weights):
         def keep(key):
@@ -113,6 +108,7 @@ class Model(nn.Module):
     def layers(self):
         return self.language_model.model.layers
 
+    @property
     def cast_predicate(self):
         def predicate(k):
             return "e_score_correction_bias" not in k

mlx_lm/models/lfm2-vl.py

@@ -0,0 +1,51 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten, tree_unflatten
+
+from . import lfm2
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    text_config: dict
+
+    def __post_init__(self):
+        self.text_config["tie_word_embeddings"] = False
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.language_model = lfm2.Model(lfm2.ModelArgs.from_dict(args.text_config))
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+        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.pop("multi_modal_projector", None)
+        return dict(tree_flatten(weights))
+
+    @property
+    def layers(self):
+        return self.language_model.model.layers
+
+    def make_cache(self):
+        return self.language_model.make_cache()

mlx_lm/models/lfm2.py

@@ -0,0 +1,302 @@
+# 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
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    max_position_embeddings: int
+    norm_eps: float
+    conv_bias: bool
+    conv_L_cache: int
+    block_dim: int
+    block_ff_dim: int
+    block_multiple_of: int
+    block_ffn_dim_multiplier: float
+    block_auto_adjust_ff_dim: bool
+    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):
+    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,
+    ):
+        seqlen = x.shape[1]
+        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,
+        dim: int,
+        ff_dim: int,
+        multiple_of: int,
+        auto_adjust_ff_dim: bool,
+        ffn_dim_multiplier: Optional[float],
+    ):
+        super().__init__()
+        if auto_adjust_ff_dim:
+            ff_dim = int(2 * ff_dim / 3)
+            if ffn_dim_multiplier is not None:
+                ff_dim = int(ffn_dim_multiplier * ff_dim)
+            ff_dim = multiple_of * ((ff_dim + multiple_of - 1) // multiple_of)
+
+        self.w1 = nn.Linear(dim, ff_dim, bias=False)
+        self.w3 = nn.Linear(dim, ff_dim, bias=False)
+        self.w2 = nn.Linear(ff_dim, dim, bias=False)
+
+    def __call__(self, x) -> mx.array:
+        return self.w2(nn.silu(self.w1(x)) * self.w3(x))
+
+
+class Lfm2DecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.is_attention_layer = layer_idx in args.full_attn_idxs
+
+        if self.is_attention_layer:
+            self.self_attn = Attention(args)
+        else:
+            self.conv = ShortConv(args, layer_idx)
+        self.feed_forward = MLP(
+            dim=args.block_dim,
+            ff_dim=args.block_ff_dim,
+            multiple_of=args.block_multiple_of,
+            auto_adjust_ff_dim=args.block_auto_adjust_ff_dim,
+            ffn_dim_multiplier=args.block_ffn_dim_multiplier,
+        )
+
+        self.operator_norm = nn.RMSNorm(args.hidden_size, eps=args.norm_eps)
+        self.ffn_norm = nn.RMSNorm(args.hidden_size, eps=args.norm_eps)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+
+        if self.is_attention_layer:
+            r = self.self_attn(self.operator_norm(x), mask=mask, cache=cache)
+        else:
+            r = self.conv(
+                self.operator_norm(x),
+                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)
+
+            sanitized_weights[name] = param
+        return sanitized_weights
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def make_cache(self):
+        return [
+            KVCache() if l.is_attention_layer else ArraysCache(size=1)
+            for l in self.layers
+        ]

mlx_lm/models/lille-130m.py

@@ -0,0 +1,154 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    block_size: int
+    layer_norm_eps: float
+    n_embd: int
+    n_head: int
+    n_kv_heads: int
+    n_layer: int
+    rope_theta: float
+    vocab_size: int
+    tie_word_embeddings: bool = True
+
+
+class Lille130mAttention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.n_head = args.n_head
+        self.n_kv_heads = args.n_kv_heads
+        self.head_dim = args.n_embd // args.n_head
+        self.scale = self.head_dim**-0.5
+
+        self.qkv_proj = nn.Linear(
+            args.n_embd, (args.n_head + 2 * args.n_kv_heads) * self.head_dim, bias=False
+        )
+        self.out_proj = nn.Linear(args.n_head * self.head_dim, args.n_embd, bias=False)
+
+        self.norm = nn.RMSNorm(args.n_embd, eps=args.layer_norm_eps)
+
+        self.rope = nn.RoPE(args.n_embd // args.n_head, True, 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
+
+        qkv = self.qkv_proj(self.norm(x))
+
+        q_size = self.n_head * self.head_dim
+        kv_size = self.n_kv_heads * self.head_dim
+
+        queries, keys, values = mx.split(qkv, [q_size, q_size + kv_size], axis=-1)
+
+        queries = queries.reshape(B, L, self.n_head, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.out_proj(output)
+
+
+class Lille130mMLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        hidden_dim = 256 * round(int(8 * args.n_embd / 3) / 256)
+
+        self.norm = nn.RMSNorm(args.n_embd, eps=args.layer_norm_eps)
+        self.gate_proj = nn.Linear(args.n_embd, hidden_dim, bias=False)
+        self.up_proj = nn.Linear(args.n_embd, hidden_dim, bias=False)
+        self.down_proj = nn.Linear(hidden_dim, args.n_embd, bias=False)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        h = self.norm(x)
+        return self.down_proj(nn.silu(self.gate_proj(h)) * self.up_proj(h))
+
+
+class Lille130Block(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.attention = Lille130mAttention(args)
+        self.feed_forward = Lille130mMLP(args)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        h = x + self.attention(x, mask, cache)
+        out = h + self.feed_forward(h)
+        return out
+
+
+class Lille130(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.tok_embeddings = nn.Embedding(args.vocab_size, args.n_embd)
+        self.layers = [Lille130Block(args=args) for _ in range(args.n_layer)]
+        self.norm = nn.RMSNorm(args.n_embd, eps=args.layer_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        h = self.tok_embeddings(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, cache=c)
+
+        return self.tok_embeddings.as_linear(self.norm(h))
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.transformer = Lille130(args)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        return self.transformer(inputs, cache=cache)
+
+    @property
+    def layers(self):
+        return self.transformer.layers
+
+    def sanitize(self, weights):
+        return {k: v for k, v in weights.items() if "rotary_emb" not in k}

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,29 +152,45 @@ 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,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
+        input_embeddings: Optional[mx.array] = None,
     ):
-        h = self.embed_tokens(inputs)
-
-        if mask is None:
-            mask = create_attention_mask(h, cache)
+        if input_embeddings is not None:
+            h = input_embeddings
+        else:
+            h = self.embed_tokens(inputs)
 
         if cache is None:
             cache = [None] * len(self.layers)
 
-        for layer, c in zip(self.layers, cache):
-            h = layer(h, mask, cache=c)
+        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, 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)
 
@@ -184,10 +207,10 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
+        input_embeddings: Optional[mx.array] = None,
     ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, cache, input_embeddings)
         if self.args.tie_word_embeddings:
             out = self.model.embed_tokens.as_linear(out)
         else:
@@ -206,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/llama4.py

@@ -1,7 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.
 
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Union
+from typing import Any, Optional, Union
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -17,7 +17,6 @@ class TextArgs(BaseModelArgs):
     attention_bias: bool
     attention_chunk_size: int
     head_dim: int
-    hidden_act: str
     hidden_size: int
     interleave_moe_layer_step: int
     intermediate_size: int
@@ -153,6 +152,7 @@ class MoE(nn.Module):
     def __init__(self, args):
         super().__init__()
         self.top_k = args.num_experts_per_tok
+        assert self.top_k == 1, "Only 1 expert per token supported"
         self.num_experts = args.num_local_experts
         self.experts = SwitchGLU(
             args.hidden_size, args.intermediate_size, self.num_experts
@@ -219,7 +219,6 @@ class LlamaModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
         h = self.embed_tokens(inputs)
@@ -242,21 +241,15 @@ class LlamaModel(nn.Module):
         token_pos = linds <= rinds
         chunk_mask = (block_pos == 0) & token_pos
 
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-        else:
-            chunk_mask &= mask
-
         if cache is None:
             cache = [None] * len(self.layers)
 
+        global_mask = create_attention_mask(h, cache[3])
+
         for idx, (layer, c) in enumerate(zip(self.layers, cache)):
             use_chunked_attention = (idx + 1) % 4 != 0
-            if use_chunked_attention:
-                local_mask = chunk_mask
-            else:
-                local_mask = mask
-            h = layer(h, local_mask, cache=c)
+            mask = chunk_mask if use_chunked_attention else global_mask
+            h = layer(h, mask, cache=c)
 
         return self.norm(h)
 
@@ -274,10 +267,9 @@ class LanguageModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, cache)
         return self.lm_head(out)
 
 
@@ -291,10 +283,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        return self.language_model(inputs, mask, cache)
+        return self.language_model(inputs, cache)
 
     def sanitize(self, weights):
         def to_remove(k):

mlx_lm/models/llama4_text.py

@@ -0,0 +1,181 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_attention_heads: int
+    num_hidden_layers: int
+    vocab_size: int
+    intermediate_size: int
+    intermediate_size_mlp: int
+    num_key_value_heads: int
+    rms_norm_eps: float
+    rope_theta: float
+    head_dim: int
+    tie_word_embeddings: bool
+    no_rope_layers: list
+    use_qk_norm: bool
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs, use_rope):
+        super().__init__()
+        self.args = args
+        self.n_heads = args.num_attention_heads
+        self.n_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(
+            args.hidden_size, self.n_heads * self.head_dim, bias=False
+        )
+        self.k_proj = nn.Linear(
+            args.hidden_size, self.n_kv_heads * self.head_dim, bias=False
+        )
+        self.v_proj = nn.Linear(
+            args.hidden_size, self.n_kv_heads * self.head_dim, bias=False
+        )
+        self.o_proj = nn.Linear(
+            self.n_heads * self.head_dim, args.hidden_size, bias=False
+        )
+        self.use_rope = use_rope
+        if use_rope:
+            self.rope = nn.RoPE(self.head_dim, traditional=True, base=args.rope_theta)
+        self.use_qk_norm = args.use_qk_norm
+        self.rms_norm_eps = args.rms_norm_eps
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        queries = queries.reshape(B, L, self.n_heads, -1)
+        keys = keys.reshape(B, L, self.n_kv_heads, -1)
+        if self.use_qk_norm:
+            queries = mx.fast.rms_norm(queries, None, self.rms_norm_eps)
+            keys = mx.fast.rms_norm(keys, None, self.rms_norm_eps)
+        queries = queries.transpose(0, 2, 1, 3)
+        keys = keys.transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if self.use_rope:
+            offset = cache.offset if cache is not None else 0
+            queries = self.rope(queries, offset=offset)
+            keys = self.rope(keys, offset=offset)
+
+        if cache is not None:
+            keys, values = cache.update_and_fetch(keys, values)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, dim, intermediate_size, activation=nn.silu):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(dim, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, 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))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs, use_rope):
+        super().__init__()
+        self.self_attn = Attention(args, use_rope)
+
+        self.feed_forward = MLP(
+            args.hidden_size,
+            args.intermediate_size_mlp,
+        )
+
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.feed_forward(self.post_attention_layernorm(h))
+        return h + r
+
+
+class LanguageModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            TransformerBlock(args=args, use_rope=args.no_rope_layers[i])
+            for i in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        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 = LanguageModel(args)
+
+        self.tie_word_embeddings = args.tie_word_embeddings
+        if not self.tie_word_embeddings:
+            self.output = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        h = self.model(inputs, cache)
+        if self.tie_word_embeddings:
+            return h @ self.model.embed_tokens.weight.T
+        else:
+            return self.output(h)
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/longcat_flash.py

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

mlx_lm/models/mamba.py

@@ -50,32 +50,6 @@ class ModelArgs(BaseModelArgs):
             self.use_bcdt_rms = True
 
 
-class DepthWiseConv1d(nn.Module):
-    def __init__(self, channels, kernel_size, bias=True, padding=0):
-        super().__init__()
-        self.channels = channels
-        self.kernel_size = kernel_size
-        self.padding = padding
-        self.weight = mx.random.normal((self.channels, kernel_size, 1))
-        self.bias = mx.zeros((channels,)) if bias else None
-
-    def __call__(self, x, cache=None):
-        B, L, C = x.shape
-        groups, K, _ = self.weight.shape
-
-        if cache is not None:
-            x = mx.concatenate([cache, x], axis=1)
-        else:
-            x = mx.pad(x, [(0, 0), (K - 1, 0), (0, 0)])
-
-        y = mx.conv_general(x, self.weight, groups=groups)
-
-        if self.bias is not None:
-            y = y + self.bias
-
-        return y, x[:, -K + 1 :, :]
-
-
 class MambaBlock(nn.Module):
     def __init__(self, args: ModelArgs):
         super().__init__()
@@ -97,11 +71,13 @@ class MambaBlock(nn.Module):
             self.hidden_size, self.intermediate_size * 2, bias=args.use_bias
         )
 
-        self.conv1d = DepthWiseConv1d(
-            channels=self.intermediate_size,
+        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=self.conv_kernel_size - 1,
+            padding=0,
         )
 
         self.x_proj = nn.Linear(
@@ -148,13 +124,15 @@ class MambaBlock(nn.Module):
         B, T, D = x.shape
         xz = self.in_proj(x)
         x, z = xz.split(indices_or_sections=2, axis=-1)
-
-        conv_out, new_conv_cache = self.conv1d(x, conv_cache)
+        K = self.conv_kernel_size
+        if conv_cache is not None:
+            x_full = mx.concatenate([conv_cache, x], axis=1)
+        else:
+            x_full = mx.pad(x, [(0, 0), (K - 1, 0), (0, 0)])
+        conv_out = self.conv1d(x_full)
+        new_conv_cache = x_full[:, -(K - 1) :, :]
         x = nn.silu(conv_out)
-
         A = -mx.exp(self.A_log)
-
-        outputs = []
         current_state = state_cache
         y = []
         for t in range(T):
@@ -228,15 +206,15 @@ class Model(nn.Module):
 
         return logits
 
-    def sanitize(self, weights):
-        for k, v in weights.items():
-            if "conv1d.weight" in k and v.shape[-1] != 1:
-                weights[k] = v.moveaxis(2, 1)
-        return weights
-
     def make_cache(self):
         return [MambaCache() for _ in range(len(self.layers))]
 
     @property
     def layers(self):
         return self.backbone.layers
+
+    def sanitize(self, weights):
+        for k, v in weights.items():
+            if "conv1d.weight" in k and v.shape[-1] != 1:
+                weights[k] = v.moveaxis(2, 1)
+        return weights

mlx_lm/models/mamba2.py

@@ -0,0 +1,245 @@
+# Copyright © 2025 Apple Inc.
+
+import math
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_ssm_mask
+from .cache import MambaCache
+from .ssm import ssm_update
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    num_heads: int
+    head_dim: int
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    state_size: int
+    num_hidden_layers: int
+    layer_norm_epsilon: float
+    conv_kernel: int
+    n_groups: int
+    use_bias: bool
+    use_conv_bias: bool
+    tie_word_embeddings: bool
+    time_step_limit: Tuple[float, float]
+    time_step_rank: Union[int, str]
+    ssm_state_size: Optional[int] = None
+    max_position_embeddings: int = 2056
+
+    def __post_init__(self):
+        if self.time_step_rank == "auto":
+            self.time_step_rank = math.ceil(self.hidden_size / 16)
+        if self.ssm_state_size is None:
+            self.ssm_state_size = self.state_size
+
+
+class MambaRMSNormGated(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = mx.ones(hidden_size)
+
+    def __call__(self, hidden_states: mx.array, gate: mx.array = None) -> mx.array:
+        if gate is not None:
+            hidden_states = hidden_states * nn.silu(gate)
+        return mx.fast.rms_norm(hidden_states, self.weight, self.eps)
+
+
+class Mamba2Block(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.num_heads = args.num_heads
+        self.hidden_size = args.hidden_size
+        self.ssm_state_size = args.ssm_state_size
+        self.conv_kernel_size = args.conv_kernel
+        self.intermediate_size = args.num_heads * args.head_dim
+        self.use_conv_bias = args.use_conv_bias
+        self.n_groups = args.n_groups
+        self.head_dim = args.head_dim
+        self.time_step_limit = args.time_step_limit
+        self.heads_per_group = self.num_heads // self.n_groups
+        self.use_bias = args.use_bias
+
+        self.conv_dim = self.intermediate_size + 2 * self.n_groups * self.ssm_state_size
+
+        self.conv1d = nn.Conv1d(
+            in_channels=self.conv_dim,
+            out_channels=self.conv_dim,
+            kernel_size=args.conv_kernel,
+            padding=0,
+            groups=self.conv_dim,
+            bias=args.use_conv_bias,
+        )
+
+        projection_size = self.intermediate_size + self.conv_dim + self.num_heads
+        self.in_proj = nn.Linear(self.hidden_size, projection_size, bias=args.use_bias)
+
+        self.dt_bias = mx.ones(self.num_heads)
+        self.A_log = mx.log(mx.arange(1, self.num_heads + 1, dtype=mx.float32))
+        self.D = mx.ones(self.num_heads)
+
+        self.norm = MambaRMSNormGated(
+            self.intermediate_size, eps=args.layer_norm_epsilon
+        )
+        self.out_proj = nn.Linear(
+            self.intermediate_size, self.hidden_size, bias=args.use_bias
+        )
+
+    def _apply_conv(
+        self, conv_input: mx.array, cache: Optional[MambaCache] = None
+    ) -> mx.array:
+        if cache is not None:
+            if cache[0] is None:
+                conv_state = mx.zeros(
+                    (conv_input.shape[0], self.conv_kernel_size - 1, self.conv_dim),
+                    dtype=conv_input.dtype,
+                )
+            else:
+                conv_state = cache[0]
+            padded_input = mx.concatenate([conv_state, conv_input], axis=1)
+            cache[0] = padded_input[:, -(self.conv_kernel_size - 1) :, :]
+        else:
+            padded_input = mx.pad(
+                conv_input, [(0, 0), (self.conv_kernel_size - 1, 0), (0, 0)]
+            )
+
+        conv_output = self.conv1d(padded_input)
+        return nn.silu(conv_output)
+
+    def _ssm(
+        self,
+        hidden_states: mx.array,
+        B: mx.array,
+        C: mx.array,
+        dt: mx.array,
+        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,
+        hidden_states: mx.array,
+        mask: Optional[mx.array],
+        cache: Optional[MambaCache] = None,
+    ) -> mx.array:
+        projected = self.in_proj(hidden_states)
+        gate, conv_input, dt = mx.split(
+            projected,
+            [self.intermediate_size, self.intermediate_size + self.conv_dim],
+            axis=-1,
+        )
+        if mask is not None:
+            conv_input = mx.where(mask[..., None], conv_input, 0)
+        conv_output = self._apply_conv(conv_input, cache)
+        hidden_states, 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, B, C, dt, state, mask=mask)
+        if cache:
+            cache[1] = state
+        y = self.norm(y, gate)
+        return self.out_proj(y)
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.mixer = Mamba2Block(args, layer_idx)
+        self.norm = nn.RMSNorm(args.hidden_size)
+
+    def __call__(
+        self, x: mx.array, mask: Optional[mx.array], cache: Optional[MambaCache] = None
+    ) -> mx.array:
+        output = self.mixer(self.norm(x), mask, cache)
+        return output + x
+
+
+class Mamba2(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.embeddings = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [ResidualBlock(args, i) for i in range(args.num_hidden_layers)]
+        self.norm_f = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon)
+
+    def __call__(
+        self, x: mx.array, cache: Optional[list[MambaCache]] = None
+    ) -> mx.array:
+        hidden = self.embeddings(x)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        mask = create_ssm_mask(hidden, cache[0])
+        for layer, c in zip(self.layers, cache):
+            hidden = layer(hidden, mask, c)
+
+        return self.norm_f(hidden)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.backbone = Mamba2(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[list[MambaCache]] = None
+    ) -> mx.array:
+        hidden = self.backbone(inputs, cache)
+
+        if self.args.tie_word_embeddings:
+            logits = self.backbone.embeddings.as_linear(hidden)
+        else:
+            logits = self.lm_head(hidden)
+        return logits
+
+    def make_cache(self, batch_size: int = 1) -> list[MambaCache]:
+        return [MambaCache() for _ in range(self.args.num_hidden_layers)]
+
+    @property
+    def layers(self):
+        return self.backbone.layers
+
+    def sanitize(self, weights):
+        for k, v in weights.items():
+            if "conv1d.weight" in k and v.shape[-1] != 1:
+                weights[k] = v.moveaxis(2, 1)
+        return weights

mlx_lm/models/mimo.py

@@ -1,3 +1,5 @@
+# Copyright © 2023-2025 Apple Inc.
+
 from dataclasses import dataclass
 from typing import Any, Dict, Optional, Union
 
@@ -135,17 +137,15 @@ class MiMoModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
         h = self.embed_tokens(inputs)
 
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-
         if cache is None:
             cache = [None] * len(self.layers)
 
+        mask = create_attention_mask(h, cache[0])
+
         for layer, c in zip(self.layers, cache):
             h = layer(h, mask, c)
 
@@ -166,10 +166,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, cache)
 
         if self.args.tie_word_embeddings:
             out = self.model.embed_tokens.as_linear(out)

mlx_lm/models/minicpm.py

@@ -7,6 +7,7 @@ import mlx.core as mx
 import mlx.nn as nn
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
 
 
 @dataclass
@@ -22,6 +23,7 @@ class ModelArgs(BaseModelArgs):
     num_key_value_heads: int
     scale_depth: float
     scale_emb: float
+    max_position_embeddings: Optional[int] = None
     rope_theta: float = 1000000.0
     rope_traditional: bool = False
     rope_scaling: Optional[Dict[str, Union[str, float]]] = None
@@ -67,17 +69,12 @@ class Attention(nn.Module):
             self.num_heads * self.head_dim, self.hidden_size, bias=False
         )
 
-        rope_scale = (
-            1 / args.rope_scaling["factor"]
-            if args.rope_scaling is not None and args.rope_scaling["type"] == "linear"
-            else 1
-        )
-
-        self.rope = nn.RoPE(
-            dims=self.head_dim,
-            traditional=args.rope_traditional,
-            base=self.rope_theta,
-            scale=rope_scale,
+        self.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            args.rope_traditional,
+            args.rope_scaling,
+            args.max_position_embeddings,
         )
 
     def __call__(
@@ -157,17 +154,15 @@ class MiniCPMModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
         h = self.embed_tokens(inputs) * self.args.scale_emb
 
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-
         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)
 
@@ -187,10 +182,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, cache)
 
         if not self.args.tie_word_embeddings:
             out = self.lm_head(out / (self.args.hidden_size / self.args.dim_model_base))

mlx_lm/models/minicpm3.py

@@ -7,7 +7,7 @@ import mlx.core as mx
 import mlx.nn as nn
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
-from .su_rope import SuScaledRotaryEmbedding
+from .rope_utils import SuScaledRoPE
 
 
 @dataclass
@@ -82,7 +82,7 @@ class Attention(nn.Module):
             bias=self.attention_bias,
         )
 
-        self.rope = SuScaledRotaryEmbedding(
+        self.rope = SuScaledRoPE(
             dims=args.qk_rope_head_dim,
             base=args.rope_theta,
             max_position_embeddings=args.max_position_embeddings,

mlx_lm/models/mistral3.py

@@ -0,0 +1,48 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten, tree_unflatten
+
+from . import llama
+from .base import BaseModelArgs
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    text_config: dict
+
+    def __post_init__(self):
+        self.text_config["tie_word_embeddings"] = False
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.language_model = llama.Model(llama.ModelArgs.from_dict(args.text_config))
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+        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.pop("multi_modal_projector", None)
+        return dict(tree_flatten(weights))
+
+    @property
+    def layers(self):
+        return self.language_model.model.layers

mlx_lm/models/mixtral.py

@@ -1,8 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.
 
-import math
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Dict, Optional, Union
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -162,17 +161,15 @@ class MixtralModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
         h = self.embed_tokens(inputs)
 
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-
         if cache is None:
             cache = [None] * len(self.layers)
 
+        mask = create_attention_mask(h, cache[0])
+
         for layer, c in zip(self.layers, cache):
             h = layer(h, mask, c)
 
@@ -190,10 +187,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, cache)
         return self.lm_head(out)
 
     def sanitize(self, weights):

mlx_lm/models/nemotron-nas.py

@@ -1,7 +1,6 @@
 # Copyright © 2025 Apple Inc.
 
 from dataclasses import dataclass, field
-from dataclasses import fields as dataclass_fields
 from typing import Any, Dict, List, Optional, Union
 
 import mlx.core as mx
@@ -226,13 +225,13 @@ class MLP(nn.Module):
         self.down_proj = nn.Linear(hidden_dim, dim, bias=args.mlp_bias)
         self.up_proj = nn.Linear(dim, hidden_dim, bias=args.mlp_bias)
 
-        try:
-            self.act_fn = _ACT2FN[args.hidden_act]
-        except KeyError:
+        self.act_fn = args.hidden_act
+        if self.act_fn not in _ACT2FN:
             raise ValueError(f"Unknown activation function: {args.hidden_act}")
 
     def __call__(self, x) -> mx.array:
-        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+        act_fn = _ACT2FN[self.act_fn]
+        return self.down_proj(act_fn(self.gate_proj(x)) * self.up_proj(x))
 
 
 class LinearSubblockReplacement(nn.Module):
@@ -334,19 +333,17 @@ class NemotronNASModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: Optional[mx.array] = None,
         cache: Optional[List[Any]] = None,
     ):
         h = self.embed_tokens(inputs)
 
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-
         if cache is None:
             cache = [None] * len(self.layers)
 
-        for i, layer in enumerate(self.layers):
-            h = layer(h, mask, cache=cache[i])
+        mask = create_attention_mask(h, cache[0])
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, cache=c)
 
         return self.norm(h)
 
@@ -366,10 +363,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask=None,
         cache=None,
     ):
-        out = self.model(inputs, mask=mask, cache=cache)
+        out = self.model(inputs, cache=cache)
         if self.args.tie_word_embeddings:
             out = self.model.embed_tokens.as_linear(out)
         else:

mlx_lm/models/nemotron.py

@@ -176,17 +176,14 @@ class NemotronModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
         h = self.embed_tokens(inputs)
 
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-
         if cache is None:
             cache = [None] * len(self.layers)
 
+        mask = create_attention_mask(h, cache[0])
         for layer, c in zip(self.layers, cache):
             h = layer(h, mask, cache=c)
 
@@ -205,10 +202,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, cache)
         if self.args.tie_word_embeddings:
             out = self.model.embed_tokens.as_linear(out)
         else:

mlx_lm/models/nemotron_h.py

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

mlx_lm/models/olmo.py

@@ -2,7 +2,7 @@
 
 import sys
 from dataclasses import dataclass
-from typing import Any, Optional, Tuple
+from typing import Any, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -124,17 +124,15 @@ class Transformer(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
         h = self.wte(inputs)
 
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-
         if cache is None:
             cache = [None] * len(self.blocks)
 
+        mask = create_attention_mask(h, cache[0])
+
         for block, c in zip(self.blocks, cache):
             h = block(h, mask, c)
 
@@ -154,10 +152,9 @@ class OlmoModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        return self.transformer(inputs, mask, cache)
+        return self.transformer(inputs, cache)
 
 
 class Model(nn.Module):
@@ -170,10 +167,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        return self.model(inputs, mask, cache)
+        return self.model(inputs, cache)
 
     @property
     def layers(self):

mlx_lm/models/olmo2.py

@@ -163,15 +163,12 @@ class LlamaModel(nn.Module):
         self,
         inputs: mx.array,
         cache=None,
-        mask=None,
     ):
         h = self.embed_tokens(inputs)
 
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-
         if cache is None:
             cache = [None] * len(self.layers)
+        mask = create_attention_mask(h, cache[0])
 
         for layer, c in zip(self.layers, cache):
             h = layer(h, mask, cache=c)
@@ -192,9 +189,8 @@ class Model(nn.Module):
         self,
         inputs: mx.array,
         cache=None,
-        mask=None,
     ):
-        out = self.model(inputs, cache, mask)
+        out = self.model(inputs, cache)
         if self.args.tie_word_embeddings:
             out = self.model.embed_tokens.as_linear(out)
         else:

mlx_lm/models/olmo3.py

@@ -0,0 +1,226 @@
+# 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 .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
+        )
+        self.is_full = args.layer_types[layer_idx] == "full_attention"
+
+        if self.is_full:
+            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

mlx_lm/models/olmoe.py

@@ -163,13 +163,11 @@ class OlmoeModel(nn.Module):
         self,
         inputs: mx.array,
         cache=None,
-        mask=None,
     ):
         h = self.embed_tokens(inputs)
-        if mask is None:
-            mask = create_attention_mask(h, cache)
         if cache is None:
             cache = [None] * len(self.layers)
+        mask = create_attention_mask(h, cache[0])
         for layer, c in zip(self.layers, cache):
             h = layer(h, mask, cache=c)
         return self.norm(h)
@@ -188,9 +186,8 @@ class Model(nn.Module):
         self,
         inputs: mx.array,
         cache=None,
-        mask=None,
     ):
-        out = self.model(inputs, cache, mask)
+        out = self.model(inputs, cache)
         if self.args.tie_word_embeddings:
             out = self.model.embed_tokens.as_linear(out)
         else:

mlx_lm/models/openelm.py

@@ -1,7 +1,7 @@
 # Copyright © 2023-2024 Apple Inc.
 
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Union
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -178,17 +178,14 @@ class OpenELMModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
         h = self.token_embeddings(inputs)
 
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-
         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, cache=c)
 
@@ -207,10 +204,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        out = self.transformer(inputs, mask, cache)
+        out = self.transformer(inputs, cache)
         if self.args.share_input_output_layers:
             out = self.transformer.token_embeddings.as_linear(out)
         else:

mlx_lm/models/phi.py

@@ -2,7 +2,6 @@
 
 import math
 from dataclasses import dataclass
-from typing import Tuple
 
 import mlx.core as mx
 import mlx.nn as nn
@@ -112,10 +111,9 @@ class PhiMLP(nn.Module):
         super().__init__()
         self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
         self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
-        self.act = nn.GELU(approx="precise")
 
     def __call__(self, x) -> mx.array:
-        return self.fc2(self.act(self.fc1(x)))
+        return self.fc2(nn.gelu_approx(self.fc1(x)))
 
 
 class PhiDecoderLayer(nn.Module):
@@ -143,14 +141,12 @@ class PhiModel(nn.Module):
             config.hidden_size, eps=config.layer_norm_eps
         )
 
-    def __call__(self, x, mask, cache):
+    def __call__(self, x, cache):
         x = self.embed_tokens(x)
 
-        if mask is None:
-            mask = create_attention_mask(x, cache)
-
         if cache is None:
             cache = [None] * len(self.layers)
+        mask = create_attention_mask(x, cache[0])
 
         for layer, c in zip(self.layers, cache):
             x = layer(x, mask, c)
@@ -168,10 +164,9 @@ class Model(nn.Module):
     def __call__(
         self,
         x: mx.array,
-        mask: mx.array = None,
         cache=None,
     ) -> mx.array:
-        y = self.model(x, mask, cache)
+        y = self.model(x, cache)
         return self.lm_head(y)
 
     @property

mlx_lm/models/phi3.py

@@ -7,7 +7,7 @@ import mlx.core as mx
 import mlx.nn as nn
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
-from .su_rope import SuScaledRotaryEmbedding
+from .rope_utils import SuScaledRoPE
 
 
 @dataclass
@@ -63,7 +63,7 @@ class Attention(nn.Module):
 
         rope_dim = int(head_dim * args.partial_rotary_factor)
         if args.rope_scaling and args.rope_scaling["type"] in ["longrope", "su"]:
-            self.rope = SuScaledRotaryEmbedding(
+            self.rope = SuScaledRoPE(
                 rope_dim,
                 base=args.rope_theta,
                 max_position_embeddings=args.max_position_embeddings,
@@ -171,17 +171,15 @@ class Phi3Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
         h = self.embed_tokens(inputs)
 
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-
         if cache is None:
             cache = [None] * len(self.layers)
 
+        mask = create_attention_mask(h, cache[0])
+
         for layer, c in zip(self.layers, cache):
             h = layer(h, mask, c)
 
@@ -200,10 +198,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, cache)
         if self.args.tie_word_embeddings:
             out = self.model.embed_tokens.as_linear(out)
         else:

mlx_lm/models/phi3small.py

@@ -258,19 +258,17 @@ class Phi3Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
         h = self.embed_tokens(inputs)
         if self.mup_embedding_multiplier:
             h = self.mup_embedding_multiplier * h
 
-        if mask is None:
-            mask = create_attention_mask(h, cache, return_array=True)
-
         if cache is None:
             cache = [None] * len(self.layers)
 
+        mask = create_attention_mask(h, cache[0], return_array=True)
+
         for layer, c in zip(self.layers, cache):
             h = layer(h, mask, c)
 
@@ -292,10 +290,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, cache)
         out = self.model.embed_tokens.as_linear(out)
         if self.mup_width_multiplier:
             out = out / self.mup_width_multiplier

mlx_lm/models/phimoe.py

@@ -7,7 +7,7 @@ import mlx.core as mx
 import mlx.nn as nn
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
-from .su_rope import SuScaledRotaryEmbedding
+from .rope_utils import SuScaledRoPE
 from .switch_layers import SwitchGLU
 
 
@@ -45,7 +45,7 @@ class Attention(nn.Module):
         self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=True)
         self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=True)
 
-        self.rope = SuScaledRotaryEmbedding(
+        self.rope = SuScaledRoPE(
             head_dim,
             base=args.rope_theta,
             max_position_embeddings=args.max_position_embeddings,
@@ -155,16 +155,13 @@ class PhiMoEModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ) -> mx.array:
         h = self.embed_tokens(inputs)
 
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-
         if cache is None:
             cache = [None] * len(self.layers)
+        mask = create_attention_mask(h, cache[0])
 
         for layer, c in zip(self.layers, cache):
             h = layer(h, mask, c)
@@ -183,10 +180,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, cache)
         return self.lm_head(out)
 
     def sanitize(self, weights):

mlx_lm/models/pixtral.py

@@ -0,0 +1,51 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten, tree_unflatten
+
+from . import llama
+from .base import BaseModelArgs
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    text_config: dict
+
+    def __post_init__(self):
+        self.text_config["tie_word_embeddings"] = False
+        self.text_config["num_attention_heads"] = self.text_config.get(
+            "num_attention_heads", 32
+        )
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.language_model = llama.Model(llama.ModelArgs.from_dict(args.text_config))
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+        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.pop("multi_modal_projector", None)
+        return dict(tree_flatten(weights))
+
+    @property
+    def layers(self):
+        return self.language_model.model.layers

mlx_lm/models/plamo.py

@@ -174,16 +174,14 @@ class PlamoModel(nn.Module):
         self,
         inputs: mx.array,
         cache: Optional[Any] = None,
-        mask: Optional[mx.array] = None,
     ) -> mx.array:
         h = self.embed_tokens(inputs)
 
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-
         if cache is None:
             cache = [None for _ in range(len(self.layers.layers))]
 
+        mask = create_attention_mask(h, cache[0])
+
         for layer, c in zip(self.layers.layers, cache):
             h = layer(h, mask, cache=c)
 
@@ -204,9 +202,8 @@ class Model(nn.Module):
         self,
         inputs: mx.array,
         cache: Optional[Any] = None,
-        mask: Optional[mx.array] = None,
     ) -> mx.array:
-        out = self.model(inputs, cache, mask)
+        out = self.model(inputs, cache)
         return self.lm_head(out)
 
     @property

mlx_lm/models/plamo2.py

@@ -7,9 +7,10 @@ from typing import Any, Optional
 import mlx.core as mx
 import mlx.nn as nn
 
-from mlx_lm.models.base import BaseModelArgs, create_attention_mask
+from mlx_lm.models.base import BaseModelArgs, create_attention_mask, create_ssm_mask
 
 from .cache import KVCache, MambaCache
+from .ssm import ssm_update
 
 
 @dataclass
@@ -53,174 +54,8 @@ class RMSNorm(nn.Module):
         )
 
 
-def get_initial_dt_bias(num_heads: int) -> mx.array:
-    dt_min = 0.001
-    dt_max = 0.1
-    dt = mx.exp(
-        mx.random.uniform(shape=(num_heads,)) * (math.log(dt_max) - math.log(dt_min))
-        + math.log(dt_min)
-    )
-    dt = mx.clip(dt, a_min=1e-4, a_max=None)
-    inv_dt = dt + mx.log(-mx.expm1(-dt))
-    return inv_dt
-
-
-def get_initial_A(num_heads: int) -> mx.array:
-    A = mx.arange(1, num_heads + 1, dtype=mx.float32)
-    return mx.log(A)
-
-
-# From: https://github.com/state-spaces/mamba/blob/0cce0fa645f100f00620ddf2333c2b7712abfdec/mamba_ssm/ops/triton/selective_state_update.py#L219
-def selective_state_update_ref(
-    state, x, dt, A, B, C, D=None, z=None, dt_bias=None, dt_softplus=False
-) -> tuple[mx.array, mx.array]:
-    """
-    Argument:
-        state: (batch, dim, dstate) or (batch, nheads, dim, dstate)
-        x: (batch, dim) or (batch, nheads, dim)
-        dt: (batch, dim) or (batch, nheads, dim)
-        A: (dim, dstate) or (nheads, dim, dstate)
-        B: (batch, dstate) or (batch, ngroups, dstate)
-        C: (batch, dstate) or (batch, ngroups, dstate)
-        D: (dim,) or (nheads, dim)
-        z: (batch, dim) or (batch, nheads, dim)
-        dt_bias: (dim,) or (nheads, dim)
-    Return:
-        out: (batch, dim) or (batch, nheads, dim)
-    """
-    has_heads = state.ndim > 3
-    if state.ndim == 3:
-        state = mx.expand_dims(state, 1)
-    if x.ndim == 2:
-        x = mx.expand_dims(x, 1)
-    if dt.ndim == 2:
-        dt = mx.expand_dims(dt, 1)
-    if A.ndim == 2:
-        A = mx.expand_dims(A, 0)
-    if B.ndim == 2:
-        B = mx.expand_dims(B, 1)
-    if C.ndim == 2:
-        C = mx.expand_dims(C, 1)
-    if D is not None and D.ndim == 1:
-        D = mx.expand_dims(D, 0)
-    if z is not None and z.ndim == 2:
-        z = mx.expand_dims(z, 1)
-    if dt_bias is not None and dt_bias.ndim == 1:
-        dt_bias = mx.expand_dims(dt_bias, 0)
-    batch, nheads, dim, dstate = state.shape
-    assert x.shape == (batch, nheads, dim)
-    assert dt.shape == x.shape
-    assert A.shape == (nheads, dim, dstate)
-    ngroups = B.shape[1]
-    assert nheads % ngroups == 0, "nheads must be divisible by ngroups"
-    assert B.shape == (batch, ngroups, dstate)
-    assert C.shape == B.shape
-    if D is not None:
-        assert D.shape == (nheads, dim)
-    if z is not None:
-        assert z.shape == x.shape
-    if dt_bias is not None:
-        assert dt_bias.shape == (nheads, dim)
-        dt = dt + dt_bias
-    dt = nn.softplus(dt) if dt_softplus else dt
-    dA = mx.exp(mx.expand_dims(dt, axis=-1) * A)  # (batch, nheads, dim, dstate)
-    B = mx.reshape(
-        mx.repeat(mx.expand_dims(B, axis=2), nheads // ngroups, 2),
-        (batch, nheads, dstate),
-    )  # (batch, nheads, dstate)
-    C = mx.reshape(
-        mx.repeat(mx.expand_dims(C, axis=2), nheads // ngroups, 2),
-        (batch, nheads, dstate),
-    )  # (batch, nheads, dstate)
-    dB = mx.expand_dims(dt, axis=-1) * mx.expand_dims(
-        B, axis=-2
-    )  # (batch, nheads, dim, dstate)
-    state = state * dA + dB * mx.expand_dims(x, axis=-1)  # (batch, dim, dstate)
-    out = mx.einsum("bhdn,bhn->bhd", state.astype(C.dtype), C)
-    if D is not None:
-        out += (x * D).astype(out.dtype)
-    out = (out if z is None else out * nn.silu(z)).astype(x.dtype)
-    if not has_heads:
-        out = out.squeeze(1)
-    return out, state
-
-
-def ssd_update_state(
-    ssm_state: mx.array,
-    x: mx.array,
-    dt: mx.array,
-    A: mx.array,
-    B: mx.array,
-    C: mx.array,
-    D: mx.array,
-    z: mx.array,
-    dt_bias: mx.array,
-    dt_softplus: bool,
-) -> tuple[mx.array, mx.array]:
-    assert ssm_state.dtype == mx.float32
-    dtype = x.dtype
-
-    hidden_size_per_head = x.shape[-1]
-    d_state = B.shape[-1]
-    A = mx.broadcast_to(
-        A[:, None, None], (A.shape[0], hidden_size_per_head, d_state)
-    ).astype(mx.float32)
-    dt = mx.broadcast_to(
-        dt[..., None], (dt.shape[0], dt.shape[1], hidden_size_per_head)
-    )
-    dt_bias = mx.broadcast_to(
-        dt_bias[:, None], (dt_bias.shape[0], hidden_size_per_head)
-    )
-    D = mx.broadcast_to(D[:, None], (D.shape[0], hidden_size_per_head))
-    out, ssm_state = selective_state_update_ref(
-        ssm_state,
-        x.astype(dtype),
-        dt.astype(dtype),
-        A.astype(mx.float32),
-        B.astype(dtype),
-        C.astype(dtype),
-        D.astype(mx.float32),
-        z.astype(dtype),
-        dt_bias.astype(mx.float32),
-        dt_softplus=dt_softplus,
-    )
-    return out[:, None], ssm_state
-
-
-def ssd_chunk_scan_combined(
-    x: mx.array,
-    dt: mx.array,
-    A: mx.array,
-    B: mx.array,
-    C: mx.array,
-    D: mx.array,
-    z: mx.array,
-    dt_bias: mx.array,
-    dt_softplus: bool,
-    ssm_state: mx.array,
-) -> tuple[mx.array, mx.array]:
-    assert ssm_state.dtype == mx.float32
-    length = x.shape[1]
-    ys = []
-    for i in range(length):
-        y, ssm_state = ssd_update_state(
-            ssm_state,
-            x[:, i],
-            dt[:, i],
-            A,
-            B[:, i],
-            C[:, i],
-            D if D.ndim == 1 else D[:, i],
-            z=z[:, i],
-            dt_bias=dt_bias,
-            dt_softplus=dt_softplus,
-        )
-        ys.append(y)
-    return mx.concatenate(ys, axis=1), ssm_state
-
-
 def causal_conv1d_update(conv_state, x, weight) -> tuple[mx.array, mx.array]:
-    _, seqlen, dim = x.shape
+    dim = x.shape[-1]
     state_len = conv_state.shape[-2]
     x = mx.concatenate([conv_state, x], axis=-2)
     conv_state = x[:, -state_len:]
@@ -229,7 +64,7 @@ def causal_conv1d_update(conv_state, x, weight) -> tuple[mx.array, mx.array]:
         weight,
         padding=0,
         groups=dim,
-    )[:, -seqlen:]
+    )
     return nn.silu(out), conv_state
 
 
@@ -265,9 +100,10 @@ class Mamba(nn.Module):
         )
         self.dt_proj = nn.Linear(self.dt_dim, self.num_heads, bias=False)
 
-        self.dt_bias = get_initial_dt_bias(self.num_heads)
-        self.A_log = get_initial_A(self.num_heads)
-        self.D = mx.ones(self.num_heads, dtype=mx.float32)
+        self.dt_bias = mx.zeros(shape=(self.num_heads,))
+        self.A_log = mx.log(mx.arange(1, self.num_heads + 1, dtype=mx.float32))
+
+        self.D = mx.ones(self.num_heads)
 
         self.dt_norm_weight = mx.ones(self.dt_dim)
         self.B_norm_weight = mx.ones(self.d_state)
@@ -275,6 +111,34 @@ class Mamba(nn.Module):
 
         self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
 
+    def _ssm(
+        self,
+        x: 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, _ = x.shape
+
+        x = x.reshape(batch_size, seq_len, self.num_heads, self.hidden_size_per_head)
+        B = B.reshape(batch_size, seq_len, 1, self.d_state)
+        C = C.reshape(batch_size, seq_len, 1, self.d_state)
+
+        y, state = ssm_update(
+            x,
+            self.A_log,
+            B,
+            C,
+            self.D,
+            dt,
+            self.dt_bias,
+            state,
+            mask=mask,
+        )
+        return y.reshape(batch_size, seq_len, self.intermediate_size), state
+
     def __call__(
         self,
         hidden_states: mx.array,
@@ -285,16 +149,11 @@ class Mamba(nn.Module):
 
         if cache is not None and cache[0] is not None:
             conv_state = cache[0]
-            ssm_state = cache[1]
         else:
             conv_state = mx.zeros(
                 (bsize, self.d_conv - 1, self.intermediate_size),
                 dtype=hidden_states.dtype,
             )
-            ssm_state = mx.zeros(
-                (bsize, self.num_heads, self.hidden_size_per_head, self.d_state),
-                dtype=mx.float32,
-            )
 
         zx = self.in_proj(hidden_states)
         zx = zx.reshape(bsize, length, self.num_heads, -1)
@@ -309,9 +168,10 @@ class Mamba(nn.Module):
         )
 
         x = x.reshape(bsize, -1, self.num_heads * self.hidden_size_per_head)
+        if mask is not None:
+            x = mx.where(mask[..., None], x, 0)
         x, conv_state = causal_conv1d_update(conv_state, x, self.conv1d.weight)
         BCdt = self.bcdt_proj(x)
-        x = x.reshape(bsize, length, self.num_heads, -1)
         B, C, dt = mx.split(BCdt, [self.d_state, self.d_state * 2], axis=-1)
 
         A = -mx.exp(self.A_log.astype(mx.float32))  # (num_heads,)
@@ -319,28 +179,21 @@ class Mamba(nn.Module):
         B = mx.fast.rms_norm(B, self.B_norm_weight, self.config.rms_norm_eps)
         C = mx.fast.rms_norm(C, self.C_norm_weight, self.config.rms_norm_eps)
 
-        # (bsize, length, num_heads, 1)
-        dt = self.dt_proj(dt)[..., None]
-
-        out, ssm_state = ssd_chunk_scan_combined(
+        # (bsize, length, num_heads)
+        dt = self.dt_proj(dt)
+        out, ssm_state = self._ssm(
             x,
-            dt.reshape(bsize, length, -1),
-            A,
             B,
             C,
-            D=self.D,
-            z=z,
-            dt_bias=self.dt_bias,
-            dt_softplus=True,
-            ssm_state=ssm_state,
+            dt,
+            cache[1] if cache else None,
+            mask,
         )
-
+        out = out * nn.silu(z.flatten(-2))
         if cache is not None:
             cache[0] = conv_state
             cache[1] = ssm_state
-        y = self.out_proj(out.reshape(bsize, length, -1))
-
-        return y
+        return self.out_proj(out)
 
 
 class Attention(nn.Module):
@@ -506,13 +359,27 @@ class PlamoDecoder(nn.Module):
             PlamoDecoderLayer(config, is_mamba=is_mamba(config, i))
             for i in range(config.num_hidden_layers)
         ]
+        self.ssm_idx = 0 if config.mamba_enabled else None
+        self.fa_idx = config.mamba_step // 2
 
-    def __call__(self, x: mx.array, mask: mx.array, cache):
-        for i, decoder_layer in enumerate(self.layers):
-            x = decoder_layer(
+    def __call__(self, x: mx.array, cache):
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        attn_mask = create_attention_mask(x, cache[self.fa_idx])
+        if self.ssm_idx is not None:
+            mamba_mask = create_ssm_mask(x, cache[self.ssm_idx])
+        else:
+            mamba_mask = None
+
+        for (
+            l,
+            c,
+        ) in zip(self.layers, cache):
+            x = l(
                 x,
-                mask=mask,
-                cache=cache[i],
+                mask=mamba_mask if l.is_mamba else attn_mask,
+                cache=c,
             )
         return x
 
@@ -531,23 +398,14 @@ class PlamoModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: Optional[mx.array] = None,
         cache=None,
     ):
         batch_size, seq_length = inputs.shape
 
         h = self.embed_tokens(inputs)
 
-        if mask is None:
-            mask = create_attention_mask(h, [cache[1]] if cache is not None else None)
-
-        if cache is None:
-            cache = [None] * len(self.layers.layers)
-
-        # decoder layers
         out = self.layers(
             h,
-            mask,
             cache,
         )
 
@@ -579,12 +437,9 @@ class Model(nn.Module):
         # full_attn = self.layer_idx in self.config.full_attention_idx
         return [MambaCache() if l.is_mamba else KVCache() for l in self.layers]
 
-    def __call__(
-        self, inputs: mx.array, mask: Optional[mx.array] = None, cache=None
-    ) -> mx.array:
+    def __call__(self, inputs: mx.array, cache=None) -> mx.array:
         outputs = self.model(
             inputs=inputs,
-            mask=None,
             cache=cache,
         )
         if self.config.tie_word_embeddings:

mlx_lm/models/qwen.py

@@ -120,14 +120,12 @@ class QwenModel(nn.Module):
         self.h = [TransformerBlock(args) for _ in range(args.num_hidden_layers)]
         self.ln_f = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon)
 
-    def __call__(self, inputs, mask=None, cache=None):
+    def __call__(self, inputs, cache=None):
         x = self.wte(inputs)
 
-        if mask is None:
-            mask = create_attention_mask(x, cache)
-
         if cache is None:
             cache = [None] * len(self.h)
+        mask = create_attention_mask(x, cache[0])
 
         for layer, c in zip(self.h, cache):
             x = layer(x, mask, c)
@@ -148,10 +146,9 @@ class Model(nn.Module):
     def __call__(
         self,
         x: mx.array,
-        mask: mx.array = None,
         cache=None,
     ) -> mx.array:
-        y = self.transformer(x, mask, cache)
+        y = self.transformer(x, cache)
         return self.lm_head(y)
 
     @property

mlx_lm/models/qwen2.py

@@ -135,16 +135,17 @@ class Qwen2Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
+        input_embeddings: Optional[mx.array] = None,
     ):
-        h = self.embed_tokens(inputs)
-
-        if mask is None:
-            mask = create_attention_mask(h, cache)
+        if input_embeddings is not None:
+            h = input_embeddings
+        else:
+            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)
@@ -164,10 +165,10 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
+        input_embeddings: Optional[mx.array] = None,
     ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, cache, input_embeddings)
         if self.args.tie_word_embeddings:
             out = self.model.embed_tokens.as_linear(out)
         else:

mlx_lm/models/qwen2_moe.py

@@ -1,6 +1,5 @@
 # Copyright © 2023-2024 Apple Inc.
 
-import math
 from dataclasses import dataclass
 from typing import Any, Dict, Optional, Union
 
@@ -187,17 +186,15 @@ class Qwen2MoeModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
         h = self.embed_tokens(inputs)
 
-        if mask is None:
-            mask = create_attention_mask(h, cache)
-
         if cache is None:
             cache = [None] * len(self.layers)
 
+        mask = create_attention_mask(h, cache[0])
+
         for layer, c in zip(self.layers, cache):
             h = layer(h, mask, c)
 
@@ -215,10 +212,9 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, mask, cache)
+        out = self.model(inputs, cache)
         return self.lm_head(out)
 
     def sanitize(self, weights):

mlx_lm/models/qwen2_vl.py

@@ -0,0 +1,58 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten, tree_unflatten
+
+from . import qwen2
+from .base import BaseModelArgs
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    text_config: dict
+
+    @classmethod
+    def from_dict(cls, params):
+        if "text_config" not in params:
+            return cls(model_type=params["model_type"], text_config=params)
+        return 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 = qwen2.Model(qwen2.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.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