Initial example

* gptq

* speedup

.circleci/config.yml

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

.gitignore

@@ -0,0 +1,139 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Vim
+*.swp
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# IDE files
+.idea/
+.vscode/
+
+# .DS_Store files
+.DS_Store

.pre-commit-config.yaml

@@ -0,0 +1,11 @@
+repos:
+-   repo: https://github.com/psf/black-pre-commit-mirror
+    rev: 25.1.0
+    hooks:
+    -   id: black
+-   repo: https://github.com/pycqa/isort
+    rev: 6.0.0
+    hooks:
+    -   id: isort
+        args:
+            - --profile=black

ACKNOWLEDGMENTS.md

@@ -5,12 +5,8 @@ with a short description of your contribution(s) below. For example:
 
 - Jane Smith: Added the `foo` example.
 
-MLX Examples was developed with contributions from the following individuals:
+MLX LM was developed with contributions from the following individuals:
 
-- Juarez Bochi: Added support for T5 models.
-- Sarthak Yadav: Added the `cifar` and `speechcommands` examples.
 - Shunta Saito: Added support for PLaMo models.
-- Gabrijel Boduljak: Implemented `CLIP`.
-- Markus Enzweiler: Added the `cvae` examples.
-- Prince Canuma: Helped add support for `Starcoder2` models.
-- Shiyu Li: Added the `Segment Anything Model`.
+- Gökdeniz Gülmez: Added support for the following architectures: OpenBMB's `MiniCPM` and `MiniCPM3`, Kyutai's `Helium`, State-Space's`Mamba v1`, Z.ai & THUKEG's `GLM4`, Rednote `dots.llm1`, Baisu's `Ernie4.5 MoE`, and Allenai's `OLMoE`; Added support for the following training algorithms: `full-fine-tuning`; Added support for the following other features: `Multiple Optimizers to choose for training`, and `reporting training metrics to WandB (Weights & Biases)`.
+- Prince Canuma: Helped add support for the following model architectures: HuggingFace's `Starcoder2`, Cohere's `Cohere (1 and 2)`, Alibaba Qwen's `Qwen (2, 3 and MoE)`, Microsoft's `Phi (3 and 3.5 MoE)`, `BitNet1.58`, Meta's `Llama (3 and 4)`, Google DeepMind's `Gemma 3`, and InterLM's `InternLM 2.5`.

CONTRIBUTING.md

@@ -1,11 +1,54 @@
 # Contributing to MLX LM 
 
+We want to make contributing to this project as easy and transparent as
+possible.
+
+## Pull Requests
+
+1. Fork and submit pull requests to the repo. 
+2. If you've added code that should be tested, add tests.
+3. Every PR should have passing tests and at least one review. 
+4. For code formatting install `pre-commit` using something like `pip install pre-commit` and run `pre-commit install`.
+   This should install hooks for running `black` and `clang-format` to ensure
+   consistent style for C++ and python code.
+ 
+   You can also run the formatters manually as follows on individual files:
+ 
+     ```bash
+     clang-format -i file.cpp
+     ```
+ 
+     ```bash
+     black file.py
+     ```
+
+     or,
+
+     ```bash
+     # single file
+     pre-commit run --files file1.py 
+
+     # specific files
+     pre-commit run --files file1.py file2.py
+     ```
+ 
+   or run `pre-commit run --all-files` to check all files in the repo.
+
+## Issues
+
+We use GitHub issues to track public bugs. Please ensure your description is
+clear and has sufficient instructions to be able to reproduce the issue.
+
+## License
+
+By contributing to mlx-lm, you agree that your contributions will be licensed
+under the LICENSE file in the root directory of this source tree.
+
+## Adding New Models
+
 Below are some tips to port LLMs available on Hugging Face to MLX.
 
-Before starting checkout the [general contribution
-guidelines](https://github.com/ml-explore/mlx-examples/blob/main/CONTRIBUTING.md).
-
-Next, from this directory, do an editable install:
+From this directory, do an editable install:
 
 ```shell
 pip install -e .
@@ -17,7 +60,7 @@ Then check if the model has weights in the
 convert it.
 
 After that, add the model file to the
-[`mlx_lm/models`](https://github.com/ml-explore/mlx-examples/tree/main/llms/mlx_lm/models)
+[`mlx_lm/models`](https://github.com/ml-explore/mlx-lm/tree/main/mlx_lm/models)
 directory. You can see other examples there. We recommend starting from a model
 that is similar to the model you are porting.
 
@@ -35,12 +78,12 @@ To determine the model layer names, we suggest either:
   in the Hugging Face repo.
 
 To add LoRA support edit
-[`mlx_lm/tuner/utils.py`](https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/tuner/utils.py#L27-L60)
+[`mlx_lm/tuner/utils.py`](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/tuner/utils.py#L27-L60)
 
 Finally, add a test for the new modle type to the [model
-tests](https://github.com/ml-explore/mlx-examples/blob/main/llms/tests/test_models.py).
+tests](https://github.com/ml-explore/mlx-lm/blob/main/tests/test_models.py).
 
-From the `llms/` directory, you can run the tests with:
+You can run the tests with:
 
 ```shell
 python -m unittest discover tests/

MANIFEST.in

@@ -1,2 +1,2 @@
-include mlx_lm/requirements.txt
+include requirements.txt
 recursive-include mlx_lm/ *.py

README.md

@@ -1,4 +1,17 @@
-## Generate Text with LLMs and MLX
+## MLX LM 
+
+MLX LM is a Python package for generating text and fine-tuning large language
+models on Apple silicon with MLX.
+
+Some key features include:
+
+* Integration with the Hugging Face Hub to easily use thousands of LLMs with a
+  single command. 
+* Support for quantizing and uploading models to the Hugging Face Hub.
+* [Low-rank and full model
+  fine-tuning](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/LORA.md)
+  with support for quantized models.
+* Distributed inference and fine-tuning with `mx.distributed`
 
 The easiest way to get started is to install the `mlx-lm` package:
 
@@ -14,11 +27,30 @@ pip install mlx-lm
 conda install -c conda-forge mlx-lm
 ```
 
-The `mlx-lm` package also has:
+### Quick Start
 
-- [LoRA and QLoRA fine-tuning](https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/LORA.md)
-- [Merging models](https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/MERGE.md)
-- [HTTP model serving](https://github.com/ml-explore/mlx-examples/blob/main/llms/mlx_lm/SERVER.md)
+To generate text with an LLM use:
+
+```bash
+mlx_lm.generate --prompt "How tall is Mt Everest?"
+```
+
+To chat with an LLM use:
+
+```bash
+mlx_lm.chat
+```
+
+This will give you a chat REPL that you can use to interact with the LLM. The
+chat context is preserved during the lifetime of the REPL.
+
+Commands in `mlx-lm` typically take command line options which let you specify
+the model, sampling parameters, and more. Use `-h` to see a list of available
+options for a command, e.g.:
+
+```bash
+mlx_lm.generate -h
+```
 
 ### Python API
 
@@ -29,7 +61,14 @@ from mlx_lm import load, generate
 
 model, tokenizer = load("mlx-community/Mistral-7B-Instruct-v0.3-4bit")
 
-response = generate(model, tokenizer, prompt="hello", verbose=True)
+prompt = "Write a story about Einstein"
+
+messages = [{"role": "user", "content": prompt}]
+prompt = tokenizer.apply_chat_template(
+    messages, add_generation_prompt=True
+)
+
+text = generate(model, tokenizer, prompt=prompt, verbose=True)
 ```
 
 To see a description of all the arguments you can do:
@@ -38,10 +77,14 @@ To see a description of all the arguments you can do:
 >>> help(generate)
 ```
 
+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.
+
 The `mlx-lm` package also comes with functionality to quantize and optionally
 upload models to the Hugging Face Hub.
 
-You can convert models in the Python API with:
+You can convert models using the Python API:
 
 ```python
 from mlx_lm import convert
@@ -64,8 +107,10 @@ To see a description of all the arguments you can do:
 
 #### Streaming
 
-For streaming generation, use the `stream_generate` function. This returns a
-generator object which streams the output text. For example,
+For streaming generation, use the `stream_generate` function. This yields
+a generation response object.
+
+For example,
 
 ```python
 from mlx_lm import load, stream_generate
@@ -75,11 +120,28 @@ model, tokenizer = load(repo)
 
 prompt = "Write a story about Einstein"
 
-for t in stream_generate(model, tokenizer, prompt, max_tokens=512):
-    print(t, end="", flush=True)
+messages = [{"role": "user", "content": prompt}]
+prompt = tokenizer.apply_chat_template(
+    messages, add_generation_prompt=True
+)
+
+for response in stream_generate(model, tokenizer, prompt, max_tokens=512):
+    print(response.text, end="", flush=True)
 print()
 ```
 
+#### Sampling
+
+The `generate` and `stream_generate` functions accept `sampler` and
+`logits_processors` keyword arguments. A sampler is any callable which accepts
+a possibly batched logits array and returns an array of sampled tokens.  The
+`logits_processors` must be a list of callables which take the token history
+and current logits as input and return the processed logits. The logits
+processors are applied in order.
+
+Some standard sampling functions and logits processors are provided in
+`mlx_lm.sample_utils`.
+
 ### Command Line
 
 You can also use `mlx-lm` from the command line with:
@@ -120,11 +182,55 @@ mlx_lm.convert \
     --upload-repo mlx-community/my-4bit-mistral
 ```
 
+Models can also be converted and quantized directly in the
+[mlx-my-repo](https://huggingface.co/spaces/mlx-community/mlx-my-repo) Hugging
+Face Space.
+
+### Long Prompts and Generations 
+
+`mlx-lm` has some tools to scale efficiently to long prompts and generations:
+
+- A rotating fixed-size key-value cache.
+- Prompt caching
+
+To use the rotating key-value cache pass the argument `--max-kv-size n` where
+`n` can be any integer. Smaller values like `512` will use very little RAM but
+result in worse quality. Larger values like `4096` or higher will use more RAM
+but have better quality.
+
+Caching prompts can substantially speedup reusing the same long context with
+different queries. To cache a prompt use `mlx_lm.cache_prompt`. For example:
+
+```bash
+cat prompt.txt | mlx_lm.cache_prompt \
+  --model mistralai/Mistral-7B-Instruct-v0.3 \
+  --prompt - \
+  --prompt-cache-file mistral_prompt.safetensors
+``` 
+
+Then use the cached prompt with `mlx_lm.generate`:
+
+```
+mlx_lm.generate \
+    --prompt-cache-file mistral_prompt.safetensors \
+    --prompt "\nSummarize the above text."
+```
+
+The cached prompt is treated as a prefix to the supplied prompt. Also notice
+when using a cached prompt, the model to use is read from the cache and need
+not be supplied explicitly.
+
+Prompt caching can also be used in the Python API in order to avoid
+recomputing the prompt. This is useful in multi-turn dialogues or across
+requests that use the same context. See the
+[example](https://github.com/ml-explore/mlx-lm/blob/main/mlx_lm/examples/chat.py)
+for more usage details.
+
 ### Supported Models
 
-The example supports Hugging Face format Mistral, Llama, and Phi-2 style
-models.  If the model you want to run is not supported, file an
-[issue](https://github.com/ml-explore/mlx-examples/issues/new) or better yet,
+`mlx-lm` supports thousands of Hugging Face format LLMs. If the model you want to
+run is not supported, file an
+[issue](https://github.com/ml-explore/mlx-lm/issues/new) or better yet,
 submit a pull request.
 
 Here are a few examples of Hugging Face models that work with this example:
@@ -140,6 +246,7 @@ Here are a few examples of Hugging Face models that work with this example:
 - [pfnet/plamo-13b-instruct](https://huggingface.co/pfnet/plamo-13b-instruct)
 - [stabilityai/stablelm-2-zephyr-1_6b](https://huggingface.co/stabilityai/stablelm-2-zephyr-1_6b)
 - [internlm/internlm2-7b](https://huggingface.co/internlm/internlm2-7b)
+- [tiiuae/falcon-mamba-7b-instruct](https://huggingface.co/tiiuae/falcon-mamba-7b-instruct)
 
 Most
 [Mistral](https://huggingface.co/models?library=transformers,safetensors&other=mistral&sort=trending),
@@ -167,3 +274,28 @@ model, tokenizer = load(
     tokenizer_config={"eos_token": "<|endoftext|>", "trust_remote_code": True},
 )
 ```
+
+### Large Models
+
+> [!NOTE]
+    This requires macOS 15.0 or higher to work.
+
+Models which are large relative to the total RAM available on the machine can
+be slow. `mlx-lm` will attempt to make them faster by wiring the memory
+occupied by the model and cache. This requires macOS 15 or higher to
+work.
+
+If you see the following warning message:
+
+> [WARNING] Generating with a model that requires ...
+
+then the model will likely be slow on the given machine. If the model fits in
+RAM then it can often be sped up by increasing the system wired memory limit.
+To increase the limit, set the following `sysctl`:
+
+```bash
+sudo sysctl iogpu.wired_limit_mb=N
+```
+
+The value `N` should be larger than the size of the model in megabytes but
+smaller than the memory size of the machine.

mlx_lm/LEARNED_QUANTS.md

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

mlx_lm/LORA.md

@@ -57,6 +57,9 @@ mlx_lm.lora \
     --iters 600
 ```
 
+To fine-tune the full model weights, add the `--fine-tune-type full` flag.
+Currently supported fine-tuning types are `lora` (default), `dora`, and `full`.
+
 The `--data` argument must specify a path to a `train.jsonl`, `valid.jsonl`
 when using `--train` and a path to a `test.jsonl` when using `--test`. For more
 details on the data format see the section on [Data](#Data).
@@ -67,12 +70,25 @@ mistralai/Mistral-7B-v0.1`.
 If `--model` points to a quantized model, then the training will use QLoRA,
 otherwise it will use regular LoRA.
 
-By default, the adapter config and weights are saved in `adapters/`. You can
-specify the output location with `--adapter-path`.
+By default, the adapter config and learned weights are saved in `adapters/`.
+You can specify the output location with `--adapter-path`.
 
 You can resume fine-tuning with an existing adapter with
 `--resume-adapter-file <path_to_adapters.safetensors>`.
 
+#### Logging
+
+You can log training metrics to Weights & Biases by passing a project name with
+the `--wandb` flag. Make sure to install wandb with `pip install wandb`.
+
+#### Prompt Masking
+
+The default training computes a loss for every token in the sample. You can
+ignore the prompt and compute loss for just the completion by passing
+`--mask-prompt`. Note this is only supported for `chat` and `completion`
+datasets. For `chat` datasets the final message in the message list is
+considered the completion. See the [dataset section](#Data) for more details. 
+
 ### Evaluate
 
 To compute test set perplexity use:
@@ -118,7 +134,7 @@ mlx_lm.fuse --model <path_to_model>
 ```
 
 This will by default load the adapters from `adapters/`, and save the fused
-model in the path `lora_fused_model/`. All of these are configurable.
+model in the path `fused_model/`. All of these are configurable.
 
 To upload a fused model, supply the `--upload-repo` and `--hf-path` arguments
 to `mlx_lm.fuse`. The latter is the repo name of the original model, which is
@@ -141,7 +157,7 @@ mlx_lm.fuse \
     --export-gguf
 ```
 
-This will save the GGUF model in `lora_fused_model/ggml-model-f16.gguf`. You
+This will save the GGUF model in `fused_model/ggml-model-f16.gguf`. You
 can specify the file name with `--gguf-path`.
 
 ## Data
@@ -151,59 +167,173 @@ Examples GitHub repo has an [example of the WikiSQL
 data](https://github.com/ml-explore/mlx-examples/tree/main/lora/data) in the
 correct format.
 
+Datasets can be specified in `*.jsonl` files locally or loaded from Hugging
+Face. 
+
+### Local Datasets
+
 For fine-tuning (`--train`), the data loader expects a `train.jsonl` and a
 `valid.jsonl` to be in the data directory. For evaluation (`--test`), the data
-loader expects a `test.jsonl` in the data directory.
+loader expects a `test.jsonl` in the data directory. 
 
-Currently, `*.jsonl` files support three data formats: `chat`,
-`completions`, and `text`. Here are three examples of these formats:
+Currently, `*.jsonl` files support `chat`, `tools`, `completions`, and `text`
+data formats. Here are examples of these formats:
 
 `chat`:
 
+```jsonl
+{"messages": [{"role": "system", "content": "You are a helpful assistant."}, {"role": "user", "content": "Hello."}, {"role": "assistant", "content": "How can I assistant you today."}]}
+```
+
+`tools`:
+
+```jsonl
+{"messages":[{"role":"user","content":"What is the weather in San Francisco?"},{"role":"assistant","tool_calls":[{"id":"call_id","type":"function","function":{"name":"get_current_weather","arguments":"{\"location\": \"San Francisco, USA\", \"format\": \"celsius\"}"}}]}],"tools":[{"type":"function","function":{"name":"get_current_weather","description":"Get the current weather","parameters":{"type":"object","properties":{"location":{"type":"string","description":"The city and country, eg. San Francisco, USA"},"format":{"type":"string","enum":["celsius","fahrenheit"]}},"required":["location","format"]}}}]}
+```
+
+<details>
+<summary>View the expanded single data tool format</summary>
+
 ```jsonl
 {
-  "messages": [
-    {
-      "role": "system",
-      "content": "You are a helpful assistant."
-    },
-    {
-      "role": "user",
-      "content": "Hello."
-    },
-    {
-      "role": "assistant",
-      "content": "How can I assistant you today."
-    }
-  ]
+    "messages": [
+        { "role": "user", "content": "What is the weather in San Francisco?" },
+        {
+            "role": "assistant",
+            "tool_calls": [
+                {
+                    "id": "call_id",
+                    "type": "function",
+                    "function": {
+                        "name": "get_current_weather",
+                        "arguments": "{\"location\": \"San Francisco, USA\", \"format\": \"celsius\"}"
+                    }
+                }
+            ]
+        }
+    ],
+    "tools": [
+        {
+            "type": "function",
+            "function": {
+                "name": "get_current_weather",
+                "description": "Get the current weather",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "location": {
+                            "type": "string",
+                            "description": "The city and country, eg. San Francisco, USA"
+                        },
+                        "format": { "type": "string", "enum": ["celsius", "fahrenheit"] }
+                    },
+                    "required": ["location", "format"]
+                }
+            }
+        }
+    ]
 }
 ```
 
+
+The format for the `arguments` field in a function varies for different models.
+Common formats include JSON strings and dictionaries. The example provided
+follows the format used by
+[OpenAI](https://platform.openai.com/docs/guides/fine-tuning/fine-tuning-examples)
+and [Mistral
+AI](https://github.com/mistralai/mistral-finetune?tab=readme-ov-file#instruct).
+A dictionary format is used in Hugging Face's [chat
+templates](https://huggingface.co/docs/transformers/main/en/chat_templating#a-complete-tool-use-example).
+Refer to the documentation for the model you are fine-tuning for more details.
+
+</details>
+
 `completions`:
 
 ```jsonl
-{
-  "prompt": "What is the capital of France?",
-  "completion": "Paris."
-}
+{"prompt": "What is the capital of France?", "completion": "Paris."}
 ```
 
+For the `completions` data format, a different key can be used for the prompt
+and completion by specifying the following in the YAML config:
+
+```yaml
+prompt_feature: "input"
+completion_feature: "output"
+```
+
+Here, `"input"` is the expected key instead of the default `"prompt"`, and
+`"output"` is the expected key instead of `"completion"`. 
+
 `text`:
 
 ```jsonl
-{
-  "text": "This is an example for the model."
-}
+{"text": "This is an example for the model."}
 ```
 
-Note, the format is automatically determined by the dataset. Note also, keys in
-each line not expected by the loader will be ignored.
+Note, the format is automatically determined by the dataset. Note also, keys
+in each line not expected by the loader will be ignored.
 
-For the `chat` and `completions` formats, Hugging Face [chat
-templates](https://huggingface.co/blog/chat-templates) are used. This applies
-the model's chat template by default. If the model does not have a chat
-template, then Hugging Face will use a default. For example, the final text in
-the `chat` example above with Hugging Face's default template becomes:
+> [!NOTE]
+> Each example in the datasets must be on a single line. Do not put more than
+> one example per line and do not split an example across multiple lines.
+
+### Hugging Face Datasets
+
+To use Hugging Face datasets, first install the `datasets` package:
+
+```
+pip install datasets
+```
+
+If the Hugging Face dataset is already in a supported format, you can specify
+it on the command line. For example, pass `--data mlx-community/wikisql` to
+train on the pre-formatted WikiwSQL data.
+
+Otherwise, provide a mapping of keys in the dataset to the features MLX LM
+expects. Use a YAML config to specify the Hugging Face dataset arguments. For
+example:
+
+```yaml
+hf_dataset:
+  path: "billsum"
+  prompt_feature: "text"
+  completion_feature: "summary"
+```
+
+- Use `prompt_feature` and `completion_feature` to specify keys for a
+  `completions` dataset. Use `text_feature` to specify the key for a `text`
+  dataset. Use `chat_feature` to specify the key for a chat dataset.
+
+- To specify the train, valid, or test splits, set the corresponding
+  `{train,valid,test}_split` argument. 
+
+You can specify a list of Hugging Face datasets with a list of records each
+with the same structure as above. For example:
+
+```yaml
+hf_dataset:
+  - path: "Open-Orca/OpenOrca"
+    train_split: "train[:90%]"
+    valid_split: "train[-10%:]"
+    prompt_feature: "question"
+    completion_feature: "response"
+  - path: "trl-lib/ultrafeedback_binarized"
+    train_split: "train[:90%]"
+    valid_split: "train[-10%:]"
+    chat_feature: "chosen"
+```
+
+- Arguments specified in `config` will be passed as keyword arguments to
+  [`datasets.load_dataset`](https://huggingface.co/docs/datasets/v2.20.0/en/package_reference/loading_methods#datasets.load_dataset).
+
+In general, for the `chat`, `tools` and `completions` formats, Hugging Face
+[chat
+templates](https://huggingface.co/docs/transformers/main/en/chat_templating)
+are used. This applies the model's chat template by default. If the model does
+not have a chat template, then Hugging Face will use a default. For example,
+the final text in the `chat` example above with Hugging Face's default template
+becomes:
 
 ```text
 <|im_start|>system
@@ -231,7 +361,7 @@ of memory. Here are some tips to reduce memory use should you need to do so:
    setting this to `2` or `1` will reduce memory consumption. This may slow
    things down a little, but will also reduce the memory use.
 
-3. Reduce the number of layers to fine-tune with `--lora-layers`. The default
+3. Reduce the number of layers to fine-tune with `--num-layers`. The default
    is `16`, so you can try `8` or `4`. This reduces the amount of memory
    needed for back propagation. It may also reduce the quality of the
    fine-tuned model if you are fine-tuning with a lot of data.
@@ -253,8 +383,8 @@ mlx_lm.lora \
     --model mistralai/Mistral-7B-v0.1 \
     --train \
     --batch-size 1 \
-    --lora-layers 4 \
-    --data wikisql
+    --num-layers 4 \
+    --data  mlx-community/wikisql
 ```
 
 The above command on an M1 Max with 32 GB runs at about 250
@@ -263,4 +393,5 @@ tokens-per-second, using the MLX Example
 data set.
 
 [^lora]: Refer to the [arXiv paper](https://arxiv.org/abs/2106.09685) for more details on LoRA.
+
 [^qlora]: Refer to the paper [QLoRA: Efficient Finetuning of Quantized LLMs](https://arxiv.org/abs/2305.14314)

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

@@ -17,7 +17,7 @@ mlx_lm.server --model <path_to_model_or_hf_repo>
 For example:
 
 ```shell
-mlx_lm.server --model mistralai/Mistral-7B-Instruct-v0.1
+mlx_lm.server --model mlx-community/Mistral-7B-Instruct-v0.3-4bit
 ```
 
 This will start a text generation server on port `8080` of the `localhost`
@@ -50,22 +50,28 @@ curl localhost:8080/v1/chat/completions \
 - `role_mapping`: (Optional) A dictionary to customize the role prefixes in
   the generated prompt. If not provided, the default mappings are used.
 
-- `stop`: (Optional) An array of strings or a single string. Thesse are
+- `stop`: (Optional) An array of strings or a single string. These are
   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`.
 
@@ -73,4 +79,65 @@ curl localhost:8080/v1/chat/completions \
   applying repetition penalty. Defaults to `20`.
 
 - `logit_bias`: (Optional) A dictionary mapping token IDs to their bias
-  values. Defaults to `None`.
+  values. Defaults to `None`.
+
+- `logprobs`: (Optional) An integer specifying the number of top tokens and
+  corresponding log probabilities to return for each output in the generated
+  sequence. If set, this can be any value between 1 and 10, inclusive.
+
+- `model`: (Optional) A string path to a local model or Hugging Face repo id.
+  If the path is local is must be relative to the directory the server was
+  started in.
+
+- `adapters`: (Optional) A string path to low-rank adapters. The path must be
+  relative to the directory the server was started in.
+
+- `draft_model`: (Optional) Specifies a smaller model to use for speculative
+  decoding. Set to `null` to unload.
+
+- `num_draft_tokens`: (Optional) The number of draft tokens the draft model
+  should predict at once. Defaults to `3`.
+
+### Response Fields
+
+- `id`: A unique identifier for the chat.
+
+- `system_fingerprint`: A unique identifier for the system.
+
+- `object`: Any of "chat.completion", "chat.completion.chunk" (for
+  streaming), or "text.completion".
+
+- `model`: The model repo or path (e.g. `"mlx-community/Llama-3.2-3B-Instruct-4bit"`).
+
+- `created`: A time-stamp for when the request was processed.
+
+- `choices`: A list of outputs. Each output is a dictionary containing the fields:
+    - `index`: The index in the list.
+    - `logprobs`: A dictionary containing the fields:
+        - `token_logprobs`: A list of the log probabilities for the generated
+          tokens.
+        - `tokens`: A list of the generated token ids.
+        - `top_logprobs`: A list of lists. Each list contains the `logprobs`
+          top tokens (if requested) with their corresponding probabilities.
+    - `finish_reason`: The reason the completion ended. This can be either of
+      `"stop"` or `"length"`.
+    - `message`: The text response from the model.
+
+- `usage`: A dictionary containing the fields:
+    - `prompt_tokens`: The number of prompt tokens processed.
+    - `completion_tokens`: The number of tokens generated.
+    - `total_tokens`: The total number of tokens, i.e. the sum of the above two fields.
+
+### List Models
+
+Use the `v1/models` endpoint to list available models:
+
+```shell
+curl localhost:8080/v1/models -H "Content-Type: application/json"
+```
+
+This will return a list of locally available models where each model in the
+list contains the following fields:
+
+- `id`: The Hugging Face repo id.
+- `created`: A time-stamp representing the model creation time.

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

@@ -1,4 +1,11 @@
 # Copyright © 2023-2024 Apple Inc.
 
-from .utils import convert, generate, load, stream_generate
-from .version import __version__
+import os
+
+from ._version import __version__
+
+os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1"
+
+from .convert import convert
+from .generate import generate, stream_generate
+from .utils import load

mlx_lm/__main__.py

@@ -0,0 +1,29 @@
+# Copyright © 2025 Apple Inc.
+
+import importlib
+import sys
+
+if __name__ == "__main__":
+    subcommands = {
+        "quant.awq",
+        "quant.dwq",
+        "quant.dynamic_quant",
+        "quant.gptq",
+        "cache_prompt",
+        "chat",
+        "convert",
+        "evaluate",
+        "fuse",
+        "generate",
+        "lora",
+        "server",
+        "manage",
+        "upload",
+    }
+    if len(sys.argv) < 2:
+        raise ValueError(f"CLI requires a subcommand in {subcommands}")
+    subcommand = sys.argv.pop(1)
+    if subcommand not in subcommands:
+        raise ValueError(f"CLI requires a subcommand in {subcommands}")
+    submodule = importlib.import_module(f"mlx_lm.{subcommand}")
+    submodule.main()

mlx_lm/_version.py

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

mlx_lm/cache_prompt.py

@@ -0,0 +1,166 @@
+# Copyright © 2024 Apple Inc.
+
+import argparse
+import json
+import sys
+import time
+
+import mlx.core as mx
+
+from .generate import generate_step
+from .models.cache import make_prompt_cache, save_prompt_cache
+from .utils import load
+
+DEFAULT_QUANTIZED_KV_START = 5000
+
+
+def setup_arg_parser():
+    """Set up and return the argument parser."""
+    parser = argparse.ArgumentParser(
+        description="Cache the state of a prompt to be reused with mlx_lm.generate"
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="mlx_model",
+        help="The path to the local model directory or Hugging Face repo.",
+    )
+    parser.add_argument(
+        "--adapter-path",
+        type=str,
+        help="Optional path for the trained adapter weights and config.",
+    )
+    parser.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        help="Enable trusting remote code for tokenizer",
+    )
+    parser.add_argument(
+        "--eos-token",
+        type=str,
+        default=None,
+        help="End of sequence token for tokenizer",
+    )
+    parser.add_argument(
+        "--ignore-chat-template",
+        action="store_true",
+        help="Use the raw prompt without the tokenizer's chat template.",
+    )
+    parser.add_argument(
+        "--use-default-chat-template",
+        action="store_true",
+        help="Use the default chat template",
+    )
+    parser.add_argument(
+        "--max-kv-size",
+        type=int,
+        default=None,
+        help="Set the maximum key-value cache size",
+    )
+    parser.add_argument(
+        "--prompt-cache-file",
+        help="The file to save the prompt cache in",
+        required=True,
+    )
+    parser.add_argument(
+        "--prompt",
+        required=True,
+        help="Message to be processed by the model ('-' reads from stdin)",
+    )
+    parser.add_argument(
+        "--kv-bits",
+        type=int,
+        help="Number of bits for KV cache quantization. "
+        "Defaults to no quantization.",
+        default=None,
+    )
+    parser.add_argument(
+        "--kv-group-size",
+        type=int,
+        help="Group size for KV cache quantization.",
+        default=64,
+    )
+    parser.add_argument(
+        "--quantized-kv-start",
+        help="When --kv-bits is set, start quantizing the KV cache "
+        "from this step onwards.",
+        type=int,
+        default=DEFAULT_QUANTIZED_KV_START,
+    )
+    return parser
+
+
+def main():
+    parser = setup_arg_parser()
+    args = parser.parse_args()
+
+    # Building tokenizer_config
+    tokenizer_config = {"trust_remote_code": True if args.trust_remote_code else None}
+    if args.eos_token is not None:
+        tokenizer_config["eos_token"] = args.eos_token
+
+    model, tokenizer = load(
+        args.model,
+        adapter_path=args.adapter_path,
+        tokenizer_config=tokenizer_config,
+    )
+
+    args.prompt = sys.stdin.read() if args.prompt == "-" else args.prompt
+
+    if args.use_default_chat_template:
+        if tokenizer.chat_template is None:
+            tokenizer.chat_template = tokenizer.default_chat_template
+
+    if not args.ignore_chat_template and tokenizer.chat_template is not None:
+        messages = [{"role": "user", "content": args.prompt}]
+        prompt = tokenizer.apply_chat_template(
+            messages, add_generation_prompt=False, continue_final_message=True
+        )
+
+    else:
+        prompt = tokenizer.encode(args.prompt)
+
+    cache = make_prompt_cache(model, args.max_kv_size)
+    y = mx.array(prompt)
+
+    # Process the prompt
+    start = time.time()
+    max_msg_len = 0
+
+    def callback(processed, total_tokens):
+        current = time.time()
+        speed = processed / (current - start)
+        msg = f"\rProcessed {processed:6d} tokens ({speed:6.2f} tok/s)"
+        nonlocal max_msg_len
+        max_msg_len = max(max_msg_len, len(msg))
+        print(msg + " " * (max_msg_len - len(msg)), end="", flush=True)
+
+    for _ in generate_step(
+        y,
+        model,
+        max_tokens=0,
+        prompt_cache=cache,
+        kv_bits=args.kv_bits,
+        kv_group_size=args.kv_group_size,
+        quantized_kv_start=args.quantized_kv_start,
+        prompt_progress_callback=callback,
+    ):
+        pass
+
+    print()
+    print(f"Peak memory: {mx.get_peak_memory() / 1e9:.3f} GB")
+
+    print("Saving...")
+    metadata = {}
+    metadata["model"] = args.model
+    metadata["chat_template"] = json.dumps(tokenizer.chat_template)
+    metadata["tokenizer_config"] = json.dumps(tokenizer_config)
+    save_prompt_cache(args.prompt_cache_file, cache, metadata)
+
+
+if __name__ == "__main__":
+    print(
+        "Calling `python -m mlx_lm.cache_prompt...` directly is deprecated."
+        " Use `mlx_lm.cache_prompt...` or `python -m mlx_lm cache_prompt ...` instead."
+    )
+    main()

mlx_lm/chat.py

@@ -0,0 +1,134 @@
+# Copyright © 2023-2024 Apple Inc.
+
+import argparse
+
+import mlx.core as mx
+
+from .generate import stream_generate
+from .models.cache import make_prompt_cache
+from .sample_utils import make_sampler
+from .utils import load
+
+DEFAULT_TEMP = 0.0
+DEFAULT_TOP_P = 1.0
+DEFAULT_XTC_PROBABILITY = 0.0
+DEFAULT_XTC_THRESHOLD = 0.0
+DEFAULT_SEED = None
+DEFAULT_MAX_TOKENS = 256
+DEFAULT_MODEL = "mlx-community/Llama-3.2-3B-Instruct-4bit"
+
+
+def setup_arg_parser():
+    """Set up and return the argument parser."""
+    parser = argparse.ArgumentParser(description="Chat with an LLM")
+    parser.add_argument(
+        "--model",
+        type=str,
+        help="The path to the local model directory or Hugging Face repo.",
+        default=DEFAULT_MODEL,
+    )
+    parser.add_argument(
+        "--adapter-path",
+        type=str,
+        help="Optional path for the trained adapter weights and config.",
+    )
+    parser.add_argument(
+        "--temp", type=float, default=DEFAULT_TEMP, help="Sampling temperature"
+    )
+    parser.add_argument(
+        "--top-p", type=float, default=DEFAULT_TOP_P, help="Sampling top-p"
+    )
+    parser.add_argument(
+        "--xtc-probability",
+        type=float,
+        default=DEFAULT_XTC_PROBABILITY,
+        help="Probability of XTC sampling to happen each next token",
+    )
+    parser.add_argument(
+        "--xtc-threshold",
+        type=float,
+        default=0.0,
+        help="Thresold the probs of each next token candidate to be sampled by XTC",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=DEFAULT_SEED,
+        help="PRNG seed",
+    )
+    parser.add_argument(
+        "--max-kv-size",
+        type=int,
+        help="Set the maximum key-value cache size",
+        default=None,
+    )
+    parser.add_argument(
+        "--max-tokens",
+        "-m",
+        type=int,
+        default=DEFAULT_MAX_TOKENS,
+        help="Maximum number of tokens to generate",
+    )
+    return parser
+
+
+def main():
+    parser = setup_arg_parser()
+    args = parser.parse_args()
+
+    if args.seed is not None:
+        mx.random.seed(args.seed)
+
+    model, tokenizer = load(
+        args.model,
+        adapter_path=args.adapter_path,
+        tokenizer_config={"trust_remote_code": True},
+    )
+
+    def print_help():
+        print("The command list:")
+        print("- 'q' to exit")
+        print("- 'r' to reset the chat")
+        print("- 'h' to display these commands")
+
+    print(f"[INFO] Starting chat session with {args.model}.")
+    print_help()
+    prompt_cache = make_prompt_cache(model, args.max_kv_size)
+    while True:
+        query = input(">> ")
+        if query == "q":
+            break
+        if query == "r":
+            prompt_cache = make_prompt_cache(model, args.max_kv_size)
+            continue
+        if query == "h":
+            print_help()
+            continue
+        messages = [{"role": "user", "content": query}]
+        prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
+        for response in stream_generate(
+            model,
+            tokenizer,
+            prompt,
+            max_tokens=args.max_tokens,
+            sampler=make_sampler(
+                args.temp,
+                args.top_p,
+                xtc_threshold=args.xtc_threshold,
+                xtc_probability=args.xtc_probability,
+                xtc_special_tokens=(
+                    tokenizer.encode("\n") + list(tokenizer.eos_token_ids)
+                ),
+            ),
+            prompt_cache=prompt_cache,
+        ):
+            print(response.text, flush=True, end="")
+        print()
+
+
+if __name__ == "__main__":
+    print(
+        "Calling `python -m mlx_lm.chat...` directly is deprecated."
+        " Use `mlx_lm.chat...` or `python -m mlx_lm chat ...` instead."
+    )
+    main()

mlx_lm/convert.py

@@ -1,8 +1,162 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import argparse
+from pathlib import Path
+from typing import Callable, Optional, Union
 
-from .utils import convert
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_map_with_path
+
+from .utils import (
+    dequantize_model,
+    fetch_from_hub,
+    get_model_path,
+    quantize_model,
+    save,
+    upload_to_hub,
+)
+
+
+def mixed_quant_predicate_builder(
+    recipe: str, model: nn.Module
+) -> Callable[[str, nn.Module, dict], Union[bool, dict]]:
+    high_bits = 6
+    group_size = 64
+
+    if recipe == "mixed_2_6":
+        low_bits = 2
+    elif recipe == "mixed_3_4":
+        low_bits = 3
+        high_bits = 4
+    elif recipe == "mixed_3_6":
+        low_bits = 3
+    elif recipe == "mixed_4_6":
+        low_bits = 4
+    else:
+        raise ValueError("Invalid quant recipe {recipe}")
+
+    down_keys = [k for k, _ in model.named_modules() if "down_proj" in k]
+    if len(down_keys) == 0:
+        raise ValueError("Model does not have expected keys for mixed quant.")
+
+    # Look for the layer index location in the path:
+    for layer_location, k in enumerate(down_keys[0].split(".")):
+        if k.isdigit():
+            break
+    num_layers = len(model.layers)
+
+    def mixed_quant_predicate(
+        path: str,
+        module: nn.Module,
+        config: dict,
+    ) -> Union[bool, dict]:
+        """Implements mixed quantization predicates with similar choices to, for example, llama.cpp's Q4_K_M.
+        Ref: https://github.com/ggerganov/llama.cpp/blob/917786f43d0f29b7c77a0c56767c0fa4df68b1c5/src/llama.cpp#L5265
+        By Alex Barron: https://gist.github.com/barronalex/84addb8078be21969f1690c1454855f3
+        """
+        index = (
+            int(path.split(".")[layer_location])
+            if len(path.split(".")) > layer_location
+            else 0
+        )
+        use_more_bits = (
+            index < num_layers // 8
+            or index >= 7 * num_layers // 8
+            or (index - num_layers // 8) % 3 == 2
+        )
+        if "v_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}
+        if "lm_head" in path:
+            return {"group_size": group_size, "bits": high_bits}
+
+        return {"group_size": group_size, "bits": low_bits}
+
+    return mixed_quant_predicate
+
+
+QUANT_RECIPES = ["mixed_2_6", "mixed_3_4", "mixed_3_6", "mixed_4_6"]
+
+MODEL_CONVERSION_DTYPES = ["float16", "bfloat16", "float32"]
+
+
+def convert(
+    hf_path: str,
+    mlx_path: str = "mlx_model",
+    quantize: bool = False,
+    q_group_size: int = 64,
+    q_bits: int = 4,
+    dtype: Optional[str] = None,
+    upload_repo: str = None,
+    revision: Optional[str] = None,
+    dequantize: bool = False,
+    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):
+        mlx_path = Path(mlx_path)
+
+    if mlx_path.exists():
+        raise ValueError(
+            f"Cannot save to the path {mlx_path} as it already exists."
+            " Please delete the file/directory or specify a new path to save to."
+        )
+
+    print("[INFO] Loading")
+    model_path, hf_path = get_model_path(hf_path, revision=revision)
+    model, config, tokenizer = fetch_from_hub(
+        model_path, lazy=True, trust_remote_code=trust_remote_code
+    )
+
+    if isinstance(quant_predicate, str):
+        quant_predicate = mixed_quant_predicate_builder(quant_predicate, model)
+
+    if dtype is None:
+        dtype = config.get("torch_dtype", None)
+    if dtype in MODEL_CONVERSION_DTYPES:
+        print("[INFO] Using dtype:", dtype)
+        dtype = getattr(mx, dtype)
+        cast_predicate = getattr(model, "cast_predicate", lambda _: True)
+
+        def set_dtype(k, v):
+            if cast_predicate(k) and mx.issubdtype(v.dtype, mx.floating):
+                return v.astype(dtype)
+            else:
+                return v
+
+        model.update(tree_map_with_path(set_dtype, model.parameters()))
+
+    if quantize and dequantize:
+        raise ValueError("Choose either quantize or dequantize, not both.")
+
+    if quantize:
+        print("[INFO] Quantizing")
+        model, config = quantize_model(
+            model, config, q_group_size, q_bits, quant_predicate=quant_predicate
+        )
+
+    if dequantize:
+        print("[INFO] Dequantizing")
+        config.pop("quantization", None)
+        config.pop("quantization_config", None)
+        model = dequantize_model(model)
+
+    save(
+        mlx_path,
+        model_path,
+        model,
+        tokenizer,
+        config,
+        hf_repo=hf_path,
+    )
+
+    if upload_repo is not None:
+        upload_to_hub(mlx_path, upload_repo)
 
 
 def configure_parser() -> argparse.ArgumentParser:
@@ -30,11 +184,18 @@ def configure_parser() -> argparse.ArgumentParser:
         "--q-bits", help="Bits per weight for quantization.", type=int, default=4
     )
     parser.add_argument(
-        "--dtype",
-        help="Type to save the parameters, ignored if -q is given.",
+        "--quant-predicate",
+        help=f"Mixed-bit quantization recipe.",
+        choices=QUANT_RECIPES,
         type=str,
-        choices=["float16", "bfloat16", "float32"],
-        default="float16",
+        required=False,
+    )
+    parser.add_argument(
+        "--dtype",
+        help="Type to save the non-quantized parameters. Defaults to config.json's `torch_dtype` or the current model weights dtype.",
+        type=str,
+        choices=MODEL_CONVERSION_DTYPES,
+        default=None,
     )
     parser.add_argument(
         "--upload-repo",
@@ -59,4 +220,8 @@ def main():
 
 
 if __name__ == "__main__":
+    print(
+        "Calling `python -m mlx_lm.convert ...` directly is deprecated."
+        " Use `mlx_lm.convert ...` or `python -m mlx_lm convert ...` instead."
+    )
     main()

mlx_lm/evaluate.py

@@ -0,0 +1,415 @@
+# Copyright © 2024 Apple Inc.
+
+"""
+Adapted from a PyTorch implementation by David Grangier
+"""
+
+import argparse
+import collections
+import copy
+import json
+import logging
+import os
+from importlib.metadata import version
+from pathlib import Path
+from typing import Any, Optional
+
+import lm_eval
+import mlx.core as mx
+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 tqdm import tqdm
+
+from .generate import stream_generate
+from .models.base import create_causal_mask
+from .models.cache import make_prompt_cache
+from .utils import common_prefix_len, load
+
+
+def _rstrip_until(s, untils):
+    """Limit a string <s> to the first occurrence of any substring in untils."""
+    l = len(s)
+    f = [s.find(u) for u in untils]
+    f = [l if x < 0 else x for x in f]
+    return s[: min(f)]
+
+
+def _pad_inputs(inputs):
+    lengths = np.array([len(x) for x in inputs])
+    maxlen = lengths.max()
+    padded = np.stack(
+        [np.pad(x, (0, maxlen - len(x))) for x in inputs],
+        axis=0,
+    )
+    return mx.array(padded), mx.array(lengths)
+
+
+def chat_template_fn(**extra_kwargs):
+    def apply_chat_template(self, chat_history, add_generation_prompt=True) -> str:
+        return self.tokenizer.apply_chat_template(
+            chat_history,
+            tokenize=False,
+            add_generation_prompt=add_generation_prompt,
+            continue_final_message=not add_generation_prompt,
+            **extra_kwargs,
+        )
+
+    return apply_chat_template
+
+
+@register_model("mlxlm")
+class MLXLM(LM):
+
+    tokenizer_name = lm_eval.models.huggingface.HFLM.tokenizer_name
+    apply_chat_template = chat_template_fn()
+
+    def __init__(
+        self,
+        path_or_hf_repo: str,
+        max_tokens: Optional[int] = None,
+        use_chat_template: Optional[bool] = None,
+    ) -> None:
+        super().__init__()
+        self._model, self.tokenizer = load(path_or_hf_repo)
+        self._max_tokens = max_tokens or self.tokenizer.model_max_length
+        self._batch_size = 8
+        self.use_chat_template = use_chat_template
+        if use_chat_template is None:
+            self.use_chat_template = self.tokenizer.chat_template is not None
+
+    def _process_prompt(self, prompt, step_size: int = 2048):
+        prompt = mx.array(prompt)[None]
+        cache = make_prompt_cache(self._model)
+        for i in range(0, prompt.shape[1], step_size):
+            logits = self._model(prompt[:, i : i + step_size], cache=cache)
+            mx.eval([c.state for c in cache])
+            mx.clear_cache()
+        logprobs = nn.log_softmax(logits[:, -1, :].astype(mx.float32))
+        return logprobs, cache
+
+    def _score_fn(self, inputs, cache: Optional[Any] = None, step_size: int = 2048):
+        inputs, lengths = _pad_inputs(inputs)
+        inputs, targets = inputs[..., :-1], inputs[..., 1:]
+
+        cache = cache or make_prompt_cache(self._model)
+        lengths += cache[0].offset
+
+        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)
+            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)
+
+            mx.eval(score, ig)
+            mx.clear_cache()
+
+            is_greedy.append(ig)
+            scores.append(score)
+
+        scores = mx.concatenate(scores, axis=1)
+        is_greedy = mx.concatenate(is_greedy, axis=1)
+
+        return scores, lengths, is_greedy
+
+    def _tokenize(self, texts):
+        return [
+            tuple(
+                self.tokenizer.encode(t, add_special_tokens=not self.use_chat_template)
+            )
+            for t in texts
+        ]
+
+    def loglikelihood(self, requests) -> list[tuple[float, bool]]:
+        """Compute log-likelihood of generating a continuation from a context.
+        Downstream tasks should attempt to use loglikelihood instead of other
+        LM calls whenever possible.
+        :param requests: list[Instance]
+            A list of Instance objects, with property `args` which returns a tuple (context, continuation).
+            `context: str`
+                Context string. Implementations of LM must be able to handle an
+                empty context string.
+            `continuation: str`
+                The continuation over which log likelihood will be calculated. If
+                there is a word boundary, the space should be in the continuation.
+                For example, context="hello" continuation=" world" is correct.
+        :return: list[tuple[float, bool]]
+            A list of pairs (logprob, isgreedy)
+            `logprob: float`
+                The log probability of `continuation`.
+            `isgreedy`:
+                Whether `continuation` would be generated by greedy sampling from `context`.
+        """
+        logging.info("Estimating loglikelihood for %d pairs." % len(requests))
+
+        group = mx.distributed.init()
+
+        # Group by common prefix
+        group_reqs = collections.defaultdict(list)
+        for idx, req in enumerate(requests):
+            group_reqs[req.args[0]].append((idx, req.args[1]))
+        questions = list(group_reqs.keys())
+        responses = []
+        indices = []
+        for v in group_reqs.values():
+            idx, resp = zip(*v)
+            indices.extend(idx)
+            responses.append(resp)
+
+        # split data accross ranks
+        questions = questions[group.rank() :: group.size()]
+        responses = responses[group.rank() :: group.size()]
+
+        long_completions = 0
+        scores, is_greedy = [], []
+        for q, rs in tqdm(zip(questions, responses), total=len(questions)):
+            prefix = self._tokenize([q])[0]
+            full_sequences = self._tokenize([q + r for r in rs])
+            max_completed_l = max(len(s) for s in full_sequences)
+
+            # compute truncation length
+            truncation = max(0, max_completed_l - self._max_tokens - 1)
+            orig_prefix_l = len(prefix)
+            prefix_l = max(len(prefix) - truncation, 0)
+            prefix = prefix[len(prefix) - prefix_l :]
+
+            # If the entire prompt got truncated ignore the question
+            if prefix_l == 0:
+                long_completions += 1
+                all_scores.extend([-float("inf")] * len(rs))
+                all_is_greedy.extend([False] * len(rs))
+                continue
+
+            # model scoring, returns num_requests x (logp, is_greedy, length).
+            logprobs, cache = self._process_prompt(prefix)
+            max_idx = mx.argmax(logprobs).item()
+
+            for s in full_sequences:
+                inputs = s[len(prefix) :]
+                # The logprobs from the last token of the prompt are
+                # for the first input token
+                scores.append(logprobs[0, inputs[0]].item())
+                is_greedy.append((inputs[0] == max_idx))
+
+                if len(inputs) == 1:
+                    continue
+                score, _, ig = self._score_fn(
+                    mx.array(inputs)[None, :], cache=copy.deepcopy(cache)
+                )
+                scores[-1] += mx.sum(score).item()
+                is_greedy[-1] &= mx.all(ig).item()
+
+        scores = mx.array(scores)
+        is_greedy = mx.array(is_greedy)
+
+        if long_completions > 0:
+            logging.info(
+                f"Prefix eliminated for {long_completions} requests with "
+                + "completion longer than context."
+            )
+
+        num_results = len(requests)
+
+        # all gather the results across groups
+        if group.size() > 1:
+            per_group = int(np.ceil(num_results / group.size()))
+            scores = mx.pad(scores, ((0, per_group - len(scores)),))
+            is_greedy = mx.pad(is_greedy, ((0, per_group - len(is_greedy))))
+            scores = mx.distributed.all_gather(scores[mx.newaxis], stream=mx.cpu)
+            is_greedy = mx.distributed.all_gather(is_greedy[mx.newaxis], stream=mx.cpu)
+            mx.eval(scores, is_greedy)
+            scores = scores.T.reshape(-1)
+            is_greedy = is_greedy.T.reshape(-1)
+
+        inv_sort = mx.argsort(mx.array(indices))
+        scores = scores[:num_results][inv_sort]
+        is_greedy = is_greedy[:num_results][inv_sort]
+        return list(zip(scores.tolist(), is_greedy.tolist()))
+
+    def loglikelihood_rolling(self, requests) -> list[float]:
+        """Compute full log-likelihood of a string, with no truncation, for perplexity computation
+        - We will use the full max context length of the model.
+        - For inputs that exceed the max context length, we divide the tokenized string into chunks of up to
+        the max context length.
+        - IMPORTANT: Each document's loglikelihood/perplexity is computed *separately*, unlike other implementations
+          which may simply concatenate multiple documents together.
+        - IMPORTANT: We maximize the amount of context for each prediction. Specifically, for inputs that we break into
+          multiple chunks, the last input will still a full-sized context.
+          Example:
+            Input tokens: [ 0 1 2 3 4 5 6 7 8 9 ]
+            Prefix: EOT
+            Max context length: 4
+            Resulting input/prediction pairs:
+                INPUT:  EOT   0   1   2
+                PRED:     0   1   2   3
+                INPUT:    3   4   5   6
+                PRED:     4   5   6   7
+                INPUT:    5   6   7   8
+                PRED:             8   9
+          Observe that:
+            1. Each token is predicted exactly once
+            2. For the last pair, we provide the full context, but only score the last two tokens
+        :param requests: list[Instance]
+            A list of Instance objects with property `args` which returns a tuple (context,).
+            string: str
+                String for which we are computing overall loglikelihood
+        :return: list[tuple[float]]
+            A list of tuples (logprob,)
+            logprob: float
+                The log probability of `context` conditioned on the EOT token.
+        """
+        logging.info(
+            "Estimating loglikelihood rolling for %d sequences." % len(requests)
+        )
+        inputs = self._tokenize([req.args[0] for req in requests])
+        all_scores = []
+        for i in tqdm(range(0, len(texts), self._batch_size)):
+            batch = texts[i : i + self._batch_size]
+            scores, lengths, _ = self._score_fn(batch)
+            mask = mx.arange(scores.shape[-1]) < lengths[:, None]
+            all_scores.extend((mask * scores).sum(axis=-1).tolist())
+
+        return all_scores
+
+    def generate_until(self, requests) -> list[str]:
+        """Generate greedily until a stopping sequence
+        :param requests: list[Instance]
+            A list of Instance objects with property `args` which returns a tuple (context, until).
+            context: str
+                Context string
+            until: [str]
+                The string sequences to generate until. These string sequences
+                may each span across multiple tokens, or may be part of one token.
+        :return: list[str]
+            A list of strings continuation
+            continuation: str
+                The generated continuation.
+        """
+        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
+        # {'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(
+                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)
+        return completions
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        "Evaluate an MLX model using lm-evaluation-harness."
+    )
+    parser.add_argument("--model", help="Model to evaluate", required=True)
+    parser.add_argument("--tasks", nargs="+", required=True)
+    parser.add_argument(
+        "--output-dir", default=".", help="Output directory for result files."
+    )
+    parser.add_argument("--batch-size", type=int, default=16, help="Batch size")
+    parser.add_argument("--num-shots", type=int, default=None, help="Number of shots")
+    parser.add_argument(
+        "--max-tokens",
+        type=int,
+        help="Maximum nunber of tokens to generate. Defaults to the model's max context length.",
+    )
+    parser.add_argument(
+        "--limit",
+        default=None,
+        help="Limit the number of examples per task.",
+        type=int,
+    )
+    parser.add_argument("--seed", type=int, default=123, help="Random seed.")
+    parser.add_argument(
+        "--fewshot-as-multiturn",
+        action="store_true",
+        help="Whether to provide the fewshot examples as a multiturn "
+        "conversation or a single user turn.",
+        default=False,
+    )
+    parser.add_argument(
+        "--apply-chat-template",
+        action=argparse.BooleanOptionalAction,
+        help="Specifies whether to apply a chat template to the prompt. If "
+        "the model has a chat template, this defaults to `True`, "
+        "otherwise `False`.",
+        default=None,
+    )
+    parser.add_argument(
+        "--chat-template-args",
+        type=json.loads,
+        help="""A JSON formatted string of arguments for the tokenizer's
+        apply_chat_template, e.g. '{"enable_thinking":false}'""",
+        default="{}",
+    )
+
+    args = parser.parse_args()
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    # Silence tokenizer warnings
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+    mx.random.seed(args.seed)
+
+    lm = MLXLM(
+        args.model,
+        max_tokens=args.max_tokens,
+        use_chat_template=args.apply_chat_template,
+    )
+    MLXLM.apply_chat_template = chat_template_fn(**args.chat_template_args)
+
+    results = lm_eval.simple_evaluate(
+        model=lm,
+        tasks=args.tasks,
+        fewshot_as_multiturn=args.fewshot_as_multiturn,
+        apply_chat_template=lm.use_chat_template,
+        num_fewshot=args.num_shots,
+        limit=args.limit,
+        random_seed=args.seed,
+        numpy_random_seed=args.seed,
+        torch_random_seed=args.seed,
+        fewshot_random_seed=args.seed,
+    )
+
+    file_keys = ["eval", args.model.replace("/", "_"), version("lm_eval")]
+    if args.num_shots is not None:
+        file_keys += [f"{args.num_shots:02d}"]
+    file_keys += args.tasks
+    filename = "_".join(file_keys)
+    if mx.distributed.init().rank() == 0:
+        output_path = output_dir / filename
+        output_path.write_text(json.dumps(results["results"], indent=4))
+        print("Results:")
+        for result in results["results"].values():
+            print(json.dumps(result, indent=4))

mlx_lm/examples/afm/README.md

@@ -0,0 +1,10 @@
+Apple Foundation Model in MLX
+=============================
+
+This example provides information about porting the AFM model to MLX-LM and
+training adapters with it or using it as any other open-weights model. It is
+paired with https://developer.apple.com/apple-intelligence/foundation-models-adapter/ that
+was published during WWDC 25 and to get the weights one needs to follow these
+instructions to download the toolkit.
+
+

mlx_lm/examples/afm/extract_quantized.py

@@ -0,0 +1,38 @@
+import argparse
+
+import mlx.core as mx
+
+from mlx_lm.convert import convert
+
+
+def mixed_quant(layer_path, layer, cfg):
+    if "embedding" in layer_path:
+        return {"group_size": 32, "bits": 8}
+    return hasattr(layer, "to_quantized")
+
+
+def main(argv):
+    parser = argparse.ArgumentParser(
+        description="Quantize the AFM according to its original quantization"
+    )
+    parser.add_argument("source", help="The mlx model containing the fp32 weights")
+    parser.add_argument("destination", help="The folder to save the quantized model to")
+    parser.add_argument("--copy-adapters", action="store_true")
+    parser.add_argument(
+        "--dtype", choices=["bfloat16", "float16", "float32"], default="float32"
+    )
+    args = parser.parse_args(argv)
+
+    convert(
+        args.source,
+        args.destination,
+        quantize=True,
+        q_group_size=128,
+        q_bits=2,
+        dtype=getattr(mx, args.dtype),
+        quant_predicate=mixed_quant,
+    )
+
+
+if __name__ == "__main__":
+    main(None)

mlx_lm/examples/afm/extract_weights.py

@@ -0,0 +1,249 @@
+import argparse
+import json
+import textwrap
+from pathlib import Path
+
+import torch
+from safetensors.torch import save_file
+from transformers import LlamaTokenizerFast
+
+
+def share_data(a, b):
+    return a.untyped_storage().data_ptr() == b.untyped_storage().data_ptr()
+
+
+def get_model_config():
+    return {
+        "model_type": "afm7",
+        "vocab_size": 153600,
+        "hidden_dim": 2048,
+        "num_layers": 56,
+        "num_kv_reuse_layers": 21,
+        "num_heads": 16,
+        "num_kv_heads": 2,
+        "hidden_dim_scale_factor": 3.25,
+        "rope_theta": 500000.0,
+    }
+
+
+def get_adapter_config():
+    return {
+        "num_layers": 56,
+        "lora_parameters": {
+            "rank": 32,
+            "scale": 0.5,
+            "dropout": 0.0,
+            "keys": [
+                "mlp.gate_proj",
+                "mlp.down_proj",
+                "mlp.up_proj",
+                "self_attn.qkv_proj",
+                "self_attn.q_proj",
+                "self_attn.out_proj",
+            ],
+        },
+    }
+
+
+def get_chat_template():
+    return textwrap.dedent(
+        """
+        {%- set default_system_message = "A conversation between a user and a helpful assistant." %}
+        {%- if messages[0]['role'] == 'system' %}
+            {%- set system_message = messages[0]['content'] %}
+            {%- set loop_messages = messages[1:] %}
+        {%- else %}
+            {%- set system_message = default_system_message %}
+            {%- set loop_messages = messages %}
+        {%- endif %}
+        {{- '<turn_start> system<n>' + system_message -}}
+        {% if tools %}
+            {{- ('<n>system tools: ' + (tools | map('tojson') | join('<n>'))) -}}
+        {% endif %}
+        {{- '<turn_end>' -}}
+        {% for message in loop_messages %}
+            {{- '<turn_start> ' + message['role'] + '<n>' + message['content'] + '<turn_end>' -}}
+        {% endfor %}
+        {% if add_generation_prompt is defined and add_generation_prompt %}
+            {% if messages[-1]['role'] != 'assistant' %}
+                {{- '<turn_start> assistant<n>' -}}
+            {% endif %}
+        {% endif %}"""
+    ).strip()
+
+
+def map_model_keys(state):
+    model_keys = {}
+    for old in state:
+        if "adapter" in old:
+            continue
+        if "kv_quantizer" in old:
+            continue
+
+        new = old
+        if new.startswith("layers."):
+            new = new[7:]
+            new = new.replace("layer_", "")
+            new = new.replace("attention.norm", "input_layernorm")
+            new = new.replace(".attention.", ".self_attn.")
+            new = new.replace("self_attn.output_transform", "self_attn.out_proj")
+            new = new.replace("feed_forward.norm", "post_attention_layernorm")
+            new = new.replace(".feed_forward.", ".mlp.")
+            new = new.replace("hidden_transform.linear_0", "gate_proj")
+            new = new.replace("hidden_transform.linear_1", "up_proj")
+            new = new.replace("mlp.output_transform", "mlp.down_proj")
+            if new.startswith("segment_0"):
+                new = new.replace("segment_0", "layers")
+                new = new.replace(".qkv_transform.", ".qkv_proj.")
+                new = new.replace(".fused_linear.", ".")
+                new = new.replace(".qk_norm.query_norm.", ".q_norm.")
+                new = new.replace(".qk_norm.key_norm.", ".k_norm.")
+            elif new.startswith("segment_1"):
+                new = new.replace("segment_1", "kv_reuse_layers")
+                new = new.replace(".q_transform.", ".q_proj.")
+                new = new.replace(".q_norm.query_norm.", ".q_norm.")
+        new = new.replace(".wrapped.", ".")
+        new = "model." + new
+        model_keys[old] = new
+    return model_keys
+
+
+def map_adapter_keys(state):
+    adapter_keys = {}
+    for old in state:
+        if "adapter" not in old:
+            continue
+
+        new = old
+        new = new[7:]
+        new = new.replace("layer_", "")
+        new = new.replace(".attention.", ".self_attn.")
+        new = new.replace("self_attn.output_transform", "self_attn.out_proj")
+        new = new.replace(".feed_forward.", ".mlp.")
+        new = new.replace("hidden_transform.linear_0", "gate_proj")
+        new = new.replace("hidden_transform.linear_1", "up_proj")
+        new = new.replace("mlp.output_transform", "mlp.down_proj")
+        if new.startswith("segment_0"):
+            new = new.replace("segment_0", "layers")
+            new = new.replace(".qkv_transform.", ".qkv_proj.")
+            new = new.replace(".fused_linear.", ".")
+        elif new.startswith("segment_1"):
+            new = new.replace("segment_1", "kv_reuse_layers")
+            new = new.replace(".q_transform.", ".q_proj.")
+
+        new = new.replace(".lora_0.b_transpose", ".b_transpose.0")
+        new = new.replace(".lora_1.b_transpose", ".b_transpose.1")
+        new = new.replace(".lora_2.b_transpose", ".b_transpose.2")
+        new = new.replace(".lora_0.a_transpose", ".a_transpose.0")
+        new = new.replace(".lora_1.a_transpose", ".a_transpose.1")
+        new = new.replace(".lora_2.a_transpose", ".a_transpose.2")
+        new = new.replace("adapters.base_adapter.b_transpose", "lora_b")
+        new = new.replace("adapters.base_adapter.a_transpose", "lora_a")
+        new = "model." + new
+        adapter_keys[old] = new
+    return adapter_keys
+
+
+def add_kv_quant_weights(new_state, old_state, dt):
+    for k, v in old_state.items():
+        if "range" not in k:
+            continue
+
+        v = v.tolist()
+        weight = "quant_key_scale" if "key_quantizer" in k else "quant_value_scale"
+        new_k = k[: k.find("kv_quantizer")]
+        new_k = new_k.replace("segment_0.layer_", "")
+        new_k = new_k.replace("attention", "self_attn")
+        new_k = "model." + new_k + weight
+        quant_scale = torch.tensor(max(v[0] / (-128), v[1] / 127), dtype=dt)
+        new_state[new_k] = quant_scale
+
+
+def cast(x, dt):
+    info = torch.finfo(dt)
+    a, b = info.min, info.max
+    return x.clip(a, b).to(dt)
+
+
+def main(argv):
+    parser = argparse.ArgumentParser(
+        description="Map the PT weights to MLX-LM safetensors"
+    )
+    parser.add_argument("source", help="The source weights in PT format")
+    parser.add_argument("tokenizer", help="The source tokenizer file")
+    parser.add_argument("destination", help="The folder to write the model weights in")
+    parser.add_argument(
+        "--dtype", choices=["bfloat16", "float16", "float32"], default="float32"
+    )
+    parser.add_argument(
+        "--adapter-dtype", choices=["bfloat16", "float16", "float32"], default="float32"
+    )
+    parser.add_argument(
+        "--force",
+        "-f",
+        action="store_true",
+        help="If set overwrite the weight files in the destination folder",
+    )
+    args = parser.parse_args(argv)
+
+    destination = Path(args.destination)
+    if not destination.exists():
+        destination.mkdir()
+    model_file = destination / "model.safetensors"
+    adapter_file = destination / "adapters.safetensors"
+    if (model_file.exists() or adapter_file.exists()) and not args.force:
+        print("Model files already exist. Delete them or use --force to overwrite them")
+        return
+
+    # Write the configuration files
+    with (destination / "config.json").open("w") as f:
+        json.dump(get_model_config(), f, indent=4)
+    with (destination / "adapter_config.json").open("w") as f:
+        json.dump(get_adapter_config(), f, indent=4)
+
+    # Pop the tied output transform
+    state = torch.load(args.source)
+    if share_data(state["embedding.weight"], state["output_transform.weight"]):
+        state.pop("output_transform.weight")
+
+    # Map the weights
+    model_keys = map_model_keys(state)
+    adapter_keys = map_adapter_keys(state)
+
+    # Make the new weight dictionaries
+    dt = getattr(torch, args.dtype)
+    adapter_dt = getattr(torch, args.adapter_dtype)
+    adapters = {
+        k_new: cast(state[k_old], adapter_dt) for k_old, k_new in adapter_keys.items()
+    }
+    model = {k_new: cast(state[k_old], dt) for k_old, k_new in model_keys.items()}
+    add_kv_quant_weights(model, state, dt)
+
+    # Save them to disk
+    save_file(model, model_file)
+    save_file(adapters, adapter_file)
+
+    # Save the tokenizer
+    tok = LlamaTokenizerFast(vocab_file=args.tokenizer)
+    tok.chat_template = get_chat_template()
+    tok.eos_token_ids = tok.convert_tokens_to_ids("<turn_end>")
+    tok.save_pretrained(str(destination))
+    with (destination / "tokenizer_config.json").open("r+") as f:
+        config = json.load(f)
+        config["tokenizer_class"] = "NewlineTokenizer"
+        f.seek(0)
+        json.dump(config, f, indent=4)
+        f.truncate()
+    with (destination / "tokenizer.json").open("r+") as f:
+        tok = json.load(f)
+        tok["decoder"]["decoders"].insert(
+            1,
+            {"type": "Replace", "pattern": {"String": "<n>"}, "content": "\n"},
+        )
+        f.seek(0)
+        json.dump(tok, f, indent=4)
+        f.truncate()
+
+
+if __name__ == "__main__":
+    main(None)

mlx_lm/examples/afm/requirements.txt

@@ -0,0 +1,3 @@
+tamm==0.1.0
+transformers
+torch

mlx_lm/examples/chat.py

@@ -0,0 +1,47 @@
+# Copyright © 2024 Apple Inc.
+
+"""
+An example of a multi-turn chat with prompt caching.
+"""
+
+from mlx_lm import generate, load
+from mlx_lm.models.cache import load_prompt_cache, make_prompt_cache, save_prompt_cache
+
+model, tokenizer = load("mlx-community/Mistral-7B-Instruct-v0.3-4bit")
+
+# Make the initial prompt cache for the model
+prompt_cache = make_prompt_cache(model)
+
+# User turn
+prompt = "Hi my name is <Name>."
+messages = [{"role": "user", "content": prompt}]
+prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
+
+# Assistant response
+response = generate(
+    model,
+    tokenizer,
+    prompt=prompt,
+    verbose=True,
+    prompt_cache=prompt_cache,
+)
+
+# User turn
+prompt = "What's my name?"
+messages = [{"role": "user", "content": prompt}]
+prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
+
+# Assistant response
+response = generate(
+    model,
+    tokenizer,
+    prompt=prompt,
+    verbose=True,
+    prompt_cache=prompt_cache,
+)
+
+# Save the prompt cache to disk to reuse it at a later time
+save_prompt_cache("mistral_prompt.safetensors", prompt_cache)
+
+# Load the prompt cache from disk
+prompt_cache = load_prompt_cache("mistral_prompt.safetensors")

mlx_lm/examples/generate_response.py

@@ -0,0 +1,33 @@
+# Copyright © 2024 Apple Inc.
+
+from mlx_lm import generate, load
+
+# Specify the checkpoint
+checkpoint = "mistralai/Mistral-7B-Instruct-v0.3"
+
+# Load the corresponding model and tokenizer
+model, tokenizer = load(path_or_hf_repo=checkpoint)
+
+# Specify the prompt and conversation history
+prompt = "Why is the sky blue?"
+conversation = [{"role": "user", "content": prompt}]
+
+# Transform the prompt into the chat template
+prompt = tokenizer.apply_chat_template(
+    conversation=conversation, add_generation_prompt=True
+)
+
+# Specify the maximum number of tokens
+max_tokens = 1_000
+
+# Specify if tokens and timing information will be printed
+verbose = True
+
+# Generate a response with the specified settings
+response = generate(
+    model=model,
+    tokenizer=tokenizer,
+    prompt=prompt,
+    max_tokens=max_tokens,
+    verbose=verbose,
+)

mlx_lm/examples/lora_config.yaml

@@ -1,16 +1,29 @@
 # The path to the local model directory or Hugging Face repo.
-model: "mlx_model"
+model: "mlx-community/Llama-3.2-1B-Instruct"
+
 # Whether or not to train (boolean)
 train: true
 
+# The fine-tuning method: "lora", "dora", or "full".
+fine_tune_type: lora
+
+# The Optimizer with its possible inputs
+optimizer: adamw
+# optimizer_config:
+#   adamw:
+#     betas: [0.9, 0.98]
+#     eps: 1e-6
+#     weight_decay: 0.05
+#     bias_correction: true
+
 # Directory with {train, valid, test}.jsonl files
-data: "/path/to/training/data"
+data: "mlx-community/WikiSQL"
 
 # The PRNG seed
 seed: 0
 
 # Number of layers to fine-tune
-lora_layers: 16
+num_layers: 16
 
 # Minibatch size.
 batch_size: 4
@@ -24,6 +37,9 @@ val_batches: 25
 # Adam learning rate.
 learning_rate: 1e-5
 
+# Whether to report the logs to WandB
+# wand: "wandb-project"
+
 # Number of training steps between loss reporting.
 steps_per_report: 10
 
@@ -51,9 +67,6 @@ max_seq_length: 2048
 # Use gradient checkpointing to reduce memory use.
 grad_checkpoint: false
 
-# Use DoRA instead of LoRA.
-use_dora: false
-
 # LoRA parameters can only be specified in a config file
 lora_parameters:
   # The layer keys to apply LoRA to.
@@ -69,3 +82,11 @@ lora_parameters:
 #  warmup: 100 # 0 for no warmup
 #  warmup_init: 1e-7 # 0 if not specified
 #  arguments: [1e-5, 1000, 1e-7] # passed to scheduler
+
+#hf_dataset:
+#  path: "billsum"
+#  train_split: "train[:1000]"
+#  valid_split: "train[-100:]"
+#  prompt_feature: "text"
+#  completion_feature: "summary"
+

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

@@ -0,0 +1,131 @@
+# Copyright © 2024 Apple Inc.
+
+"""
+Run with:
+
+```
+mlx.launch \
+ --hostfile /path/to/hosts.json \
+ /path/to/pipeline_generate.py \
+ --prompt "hello world"
+```
+
+Make sure you can run MLX over MPI on two hosts. For more information see the
+documentation:
+
+https://ml-explore.github.io/mlx/build/html/usage/distributed.html).
+"""
+
+import argparse
+import json
+import resource
+from pathlib import Path
+
+import mlx.core as mx
+from huggingface_hub import snapshot_download
+from mlx.utils import tree_flatten
+
+from mlx_lm import load, stream_generate
+from mlx_lm.utils import load_model, load_tokenizer
+
+# Needed for 8 bit model
+resource.setrlimit(resource.RLIMIT_NOFILE, (2048, 4096))
+
+
+def download(repo: str, allow_patterns: list[str]) -> Path:
+    return Path(
+        snapshot_download(
+            repo,
+            allow_patterns=allow_patterns,
+        )
+    )
+
+
+def shard_and_load(repo):
+    # Get model path with everything but weight safetensors
+    model_path = download(
+        args.model,
+        allow_patterns=["*.json", "*.py", "tokenizer.model", "*.tiktoken", "*.txt"],
+    )
+
+    # Lazy load and shard model to figure out
+    # which weights we need
+    model, _ = load_model(model_path, lazy=True, strict=False)
+
+    group = mx.distributed.init()
+    rank = group.rank()
+    model.model.pipeline(group)
+
+    # Figure out which files we need for the local shard
+    with open(model_path / "model.safetensors.index.json", "r") as fid:
+        weight_index = json.load(fid)["weight_map"]
+
+    local_files = set()
+    for k, _ in tree_flatten(model.parameters()):
+        local_files.add(weight_index[k])
+
+    # Download weights for local shard
+    download(args.model, allow_patterns=local_files)
+
+    # Load and shard the model, and load the weights
+    tokenizer = load_tokenizer(model_path)
+    model, _ = load_model(model_path, lazy=True, strict=False)
+    model.model.pipeline(group)
+    mx.eval(model.parameters())
+
+    # Synchronize processes before generation to avoid timeout if downloading
+    # model for the first time.
+    mx.eval(mx.distributed.all_sum(mx.array(1.0), stream=mx.cpu))
+    return model, tokenizer
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="LLM pipelined inference example")
+    parser.add_argument(
+        "--model",
+        default="mlx-community/DeepSeek-R1-3bit",
+        help="HF repo or path to local model.",
+    )
+    parser.add_argument(
+        "--prompt",
+        "-p",
+        default="Write a quicksort in C++.",
+        help="Message to be processed by the model ('-' reads from stdin)",
+    )
+    parser.add_argument(
+        "--max-tokens",
+        "-m",
+        type=int,
+        default=256,
+        help="Maximum number of tokens to generate",
+    )
+    args = parser.parse_args()
+
+    group = mx.distributed.init()
+    rank = group.rank()
+
+    def rprint(*args, **kwargs):
+        if rank == 0:
+            print(*args, **kwargs)
+
+    model, tokenizer = shard_and_load(args.model)
+
+    messages = [{"role": "user", "content": args.prompt}]
+    prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
+
+    for response in stream_generate(
+        model, tokenizer, prompt, max_tokens=args.max_tokens
+    ):
+        rprint(response.text, end="", flush=True)
+
+    rprint()
+    rprint("=" * 10)
+    rprint(
+        f"Prompt: {response.prompt_tokens} tokens, "
+        f"{response.prompt_tps:.3f} tokens-per-sec"
+    )
+    rprint(
+        f"Generation: {response.generation_tokens} tokens, "
+        f"{response.generation_tps:.3f} tokens-per-sec"
+    )
+    rprint(f"Peak memory: {response.peak_memory:.3f} GB")

mlx_lm/examples/tool_use.py

@@ -0,0 +1,73 @@
+# Copyright © 2025 Apple Inc.
+
+import json
+
+from mlx_lm import generate, load
+from mlx_lm.models.cache import make_prompt_cache
+
+# Specify the checkpoint
+checkpoint = "mlx-community/Qwen2.5-32B-Instruct-4bit"
+
+# Load the corresponding model and tokenizer
+model, tokenizer = load(path_or_hf_repo=checkpoint)
+
+
+# An example tool, make sure to include a docstring and type hints
+def multiply(a: float, b: float):
+    """
+    A function that multiplies two numbers
+
+    Args:
+        a: The first number to multiply
+        b: The second number to multiply
+    """
+    return a * b
+
+
+tools = {"multiply": multiply}
+
+# Specify the prompt and conversation history
+prompt = "Multiply 12234585 and 48838483920."
+messages = [{"role": "user", "content": prompt}]
+
+prompt = tokenizer.apply_chat_template(
+    messages, add_generation_prompt=True, tools=list(tools.values())
+)
+
+prompt_cache = make_prompt_cache(model)
+
+# Generate the initial tool call:
+response = generate(
+    model=model,
+    tokenizer=tokenizer,
+    prompt=prompt,
+    max_tokens=2048,
+    verbose=True,
+    prompt_cache=prompt_cache,
+)
+
+# Parse the tool call:
+# (Note, the tool call format is model specific)
+tool_open = "<tool_call>"
+tool_close = "</tool_call>"
+start_tool = response.find(tool_open) + len(tool_open)
+end_tool = response.find(tool_close)
+tool_call = json.loads(response[start_tool:end_tool].strip())
+tool_result = tools[tool_call["name"]](**tool_call["arguments"])
+
+# Put the tool result in the prompt
+messages = [{"role": "tool", "name": tool_call["name"], "content": tool_result}]
+prompt = tokenizer.apply_chat_template(
+    messages,
+    add_generation_prompt=True,
+)
+
+# Generate the final response:
+response = generate(
+    model=model,
+    tokenizer=tokenizer,
+    prompt=prompt,
+    max_tokens=2048,
+    verbose=True,
+    prompt_cache=prompt_cache,
+)

mlx_lm/fuse.py

@@ -1,19 +1,14 @@
 import argparse
-import glob
-import shutil
 from pathlib import Path
 
 from mlx.utils import tree_flatten, tree_unflatten
 
 from .gguf import convert_to_gguf
-from .tuner.dora import DoRALinear
-from .tuner.lora import LoRALinear, LoRASwitchLinear
-from .tuner.utils import apply_lora_layers, dequantize
+from .tuner.utils import dequantize, load_adapters
 from .utils import (
     fetch_from_hub,
     get_model_path,
-    save_config,
-    save_weights,
+    save,
     upload_to_hub,
 )
 
@@ -29,7 +24,7 @@ def parse_arguments() -> argparse.Namespace:
     )
     parser.add_argument(
         "--save-path",
-        default="lora_fused_model",
+        default="fused_model",
         help="The path to save the fused model.",
     )
     parser.add_argument(
@@ -38,12 +33,6 @@ def parse_arguments() -> argparse.Namespace:
         default="adapters",
         help="Path to the trained adapter weights and config.",
     )
-    parser.add_argument(
-        "--hf-path",
-        type=str,
-        default=None,
-        help="Path to the original Hugging Face model. Required for upload if --model is a local directory.",
-    )
     parser.add_argument(
         "--upload-repo",
         help="The Hugging Face repo to upload the model to.",
@@ -73,40 +62,36 @@ 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 = apply_lora_layers(model, args.adapter_path)
+    model = load_adapters(model, args.adapter_path)
 
     fused_linears = [
-        (n, m.to_linear())
+        (n, m.fuse(de_quantize=args.de_quantize))
         for n, m in model.named_modules()
-        if isinstance(m, (LoRASwitchLinear, LoRALinear, DoRALinear))
+        if hasattr(m, "fuse")
     ]
 
-    model.update_modules(tree_unflatten(fused_linears))
+    if fused_linears:
+        model.update_modules(tree_unflatten(fused_linears))
 
     if args.de_quantize:
         print("De-quantizing model")
         model = dequantize(model)
-
-    weights = dict(tree_flatten(model.parameters()))
-
-    save_path = Path(args.save_path)
-
-    save_weights(save_path, weights)
-
-    py_files = glob.glob(str(model_path / "*.py"))
-    for file in py_files:
-        shutil.copy(file, save_path)
-
-    tokenizer.save_pretrained(save_path)
-
-    if args.de_quantize:
         config.pop("quantization", None)
 
-    save_config(config, config_path=save_path / "config.json")
+    save_path = Path(args.save_path)
+    save(
+        save_path,
+        model_path,
+        model,
+        tokenizer,
+        config,
+        hf_repo=hf_path,
+        donate_model=False,
+    )
 
     if args.export_gguf:
         model_type = config["model_type"]
@@ -114,18 +99,20 @@ def main() -> None:
             raise ValueError(
                 f"Model type {model_type} not supported for GGUF conversion."
             )
+        weights = dict(tree_flatten(model.parameters()))
         convert_to_gguf(model_path, weights, config, str(save_path / args.gguf_path))
 
     if args.upload_repo is not None:
-        hf_path = args.hf_path or (
-            args.model if not Path(args.model).exists() else None
-        )
         if hf_path is None:
             raise ValueError(
                 "Must provide original Hugging Face repo to upload local model."
             )
-        upload_to_hub(args.save_path, args.upload_repo, hf_path)
+        upload_to_hub(args.save_path, args.upload_repo)
 
 
 if __name__ == "__main__":
+    print(
+        "Calling `python -m mlx_lm.fuse...` directly is deprecated."
+        " Use `mlx_lm.fuse...` or `python -m mlx_lm fuse ...` instead."
+    )
     main()

mlx_lm/generate.py

@@ -1,17 +1,52 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import argparse
+import contextlib
+import functools
+import json
+import sys
+import time
+from dataclasses import dataclass
+from typing import (
+    Any,
+    Callable,
+    Generator,
+    List,
+    Optional,
+    Tuple,
+    Union,
+)
 
 import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_reduce
+from transformers import PreTrainedTokenizer
 
-from .utils import generate, load
+from .models import cache
+from .models.cache import (
+    QuantizedKVCache,
+    load_prompt_cache,
+)
+from .sample_utils import make_sampler
+from .tokenizer_utils import TokenizerWrapper
+from .utils import does_model_support_input_embeddings, load
 
-DEFAULT_MODEL_PATH = "mlx_model"
 DEFAULT_PROMPT = "hello"
 DEFAULT_MAX_TOKENS = 100
-DEFAULT_TEMP = 0.6
+DEFAULT_TEMP = 0.0
 DEFAULT_TOP_P = 1.0
-DEFAULT_SEED = 0
+DEFAULT_MIN_P = 0.0
+DEFAULT_TOP_K = 0
+DEFAULT_XTC_PROBABILITY = 0.0
+DEFAULT_XTC_THRESHOLD = 0.0
+DEFAULT_MIN_TOKENS_TO_KEEP = 1
+DEFAULT_SEED = None
+DEFAULT_MODEL = "mlx-community/Llama-3.2-3B-Instruct-4bit"
+DEFAULT_QUANTIZED_KV_START = 5000
+
+
+def str2bool(string):
+    return string.lower() not in ["false", "f"]
 
 
 def setup_arg_parser():
@@ -20,8 +55,11 @@ def setup_arg_parser():
     parser.add_argument(
         "--model",
         type=str,
-        default="mlx_model",
-        help="The path to the local model directory or Hugging Face repo.",
+        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(
         "--adapter-path",
@@ -29,18 +67,27 @@ def setup_arg_parser():
         help="Optional path for the trained adapter weights and config.",
     )
     parser.add_argument(
-        "--trust-remote-code",
-        action="store_true",
-        help="Enable trusting remote code for tokenizer",
-    )
-    parser.add_argument(
-        "--eos-token",
+        "--extra-eos-token",
         type=str,
-        default=None,
-        help="End of sequence token for tokenizer",
+        default=(),
+        nargs="+",
+        help="Add tokens in the list of eos tokens that stop generation.",
     )
     parser.add_argument(
-        "--prompt", default=DEFAULT_PROMPT, help="Message to be processed by the model"
+        "--system-prompt",
+        default=None,
+        help="System prompt to be used for the chat template",
+    )
+    parser.add_argument(
+        "--prompt",
+        "-p",
+        default=DEFAULT_PROMPT,
+        help="Message to be processed by the model ('-' reads from stdin)",
+    )
+    parser.add_argument(
+        "--prefill-response",
+        default=None,
+        help="Prefill response to be used for the chat template",
     )
     parser.add_argument(
         "--max-tokens",
@@ -55,7 +102,36 @@ def setup_arg_parser():
     parser.add_argument(
         "--top-p", type=float, default=DEFAULT_TOP_P, help="Sampling top-p"
     )
-    parser.add_argument("--seed", type=int, default=DEFAULT_SEED, help="PRNG seed")
+    parser.add_argument(
+        "--min-p", type=float, default=DEFAULT_MIN_P, help="Sampling min-p"
+    )
+    parser.add_argument(
+        "--top-k", type=int, default=DEFAULT_TOP_K, help="Sampling top-k"
+    )
+    parser.add_argument(
+        "--xtc-probability",
+        type=float,
+        default=DEFAULT_XTC_PROBABILITY,
+        help="Probability of XTC sampling to happen each next token",
+    )
+    parser.add_argument(
+        "--xtc-threshold",
+        type=float,
+        default=0.0,
+        help="Thresold the probs of each next token candidate to be sampled by XTC",
+    )
+    parser.add_argument(
+        "--min-tokens-to-keep",
+        type=int,
+        default=DEFAULT_MIN_TOKENS_TO_KEEP,
+        help="Minimum tokens to keep for min-p sampling.",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=DEFAULT_SEED,
+        help="PRNG seed",
+    )
     parser.add_argument(
         "--ignore-chat-template",
         action="store_true",
@@ -67,95 +143,760 @@ def setup_arg_parser():
         help="Use the default chat template",
     )
     parser.add_argument(
-        "--colorize",
-        action="store_true",
-        help="Colorize output based on T[0] probability",
+        "--chat-template-config",
+        help="Additional config for `apply_chat_template`. Should be a dictionary of"
+        " string keys to values represented as a JSON decodable string.",
+        default=None,
     )
     parser.add_argument(
-        "--cache-limit-gb",
+        "--verbose",
+        type=str2bool,
+        default=True,
+        help="Log verbose output when 'True' or 'T' or only print the response when 'False' or 'F'",
+    )
+    parser.add_argument(
+        "--max-kv-size",
         type=int,
+        help="Set the maximum key-value cache size",
         default=None,
-        help="Set the MLX cache limit in GB",
-        required=False,
+    )
+    parser.add_argument(
+        "--prompt-cache-file",
+        type=str,
+        default=None,
+        help="A file containing saved KV caches to avoid recomputing them",
+    )
+    parser.add_argument(
+        "--kv-bits",
+        type=int,
+        help="Number of bits for KV cache quantization. "
+        "Defaults to no quantization.",
+        default=None,
+    )
+    parser.add_argument(
+        "--kv-group-size",
+        type=int,
+        help="Group size for KV cache quantization.",
+        default=64,
+    )
+    parser.add_argument(
+        "--quantized-kv-start",
+        help="When --kv-bits is set, start quantizing the KV cache "
+        "from this step onwards.",
+        type=int,
+        default=DEFAULT_QUANTIZED_KV_START,
+    )
+    parser.add_argument(
+        "--draft-model",
+        type=str,
+        help="A model to be used for speculative decoding.",
+        default=None,
+    )
+    parser.add_argument(
+        "--num-draft-tokens",
+        type=int,
+        help="Number of tokens to draft when using speculative decoding.",
+        default=3,
     )
     return parser
 
 
-def colorprint(color, s):
-    color_codes = {
-        "black": 30,
-        "red": 31,
-        "green": 32,
-        "yellow": 33,
-        "blue": 34,
-        "magenta": 35,
-        "cyan": 36,
-        "white": 39,
-    }
-    ccode = color_codes.get(color, 30)
-    print(f"\033[1m\033[{ccode}m{s}\033[0m", end="", flush=True)
+# A stream on the default device just for generation
+generation_stream = mx.new_stream(mx.default_device())
 
 
-def colorprint_by_t0(s, t0):
-    if t0 > 0.95:
-        color = "white"
-    elif t0 > 0.70:
-        color = "green"
-    elif t0 > 0.30:
-        color = "yellow"
+@contextlib.contextmanager
+def wired_limit(model: nn.Module, streams: Optional[List[mx.Stream]] = None):
+    """
+    A context manager to temporarily change the wired limit.
+
+    Note, the wired limit should not be changed during an async eval.  If an
+    async eval could be running pass in the streams to synchronize with prior
+    to exiting the context manager.
+    """
+    if not mx.metal.is_available():
+        try:
+            yield
+        finally:
+            return
+
+    model_bytes = tree_reduce(
+        lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc, model, 0
+    )
+    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
+class GenerationResponse:
+    """
+    The output of :func:`stream_generate`.
+
+    Args:
+        text (str): The next segment of decoded text. This can be an empty string.
+        token (int): The next token.
+        from_draft (bool): Whether the token was generated by the draft model.
+        logprobs (mx.array): A vector of log probabilities.
+        prompt_tokens (int): The number of tokens in the prompt.
+        prompt_tps (float): The prompt processing tokens-per-second.
+        generation_tokens (int): The number of generated tokens.
+        generation_tps (float): The tokens-per-second for generation.
+        peak_memory (float): The peak memory used so far in GB.
+        finish_reason (str): The reason the response is being sent: "length", "stop" or `None`
+    """
+
+    text: str
+    token: int
+    logprobs: mx.array
+    from_draft: bool
+    prompt_tokens: int
+    prompt_tps: float
+    generation_tokens: int
+    generation_tps: float
+    peak_memory: float
+    finish_reason: Optional[str] = None
+
+
+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
+                )
+
+
+def generate_step(
+    prompt: mx.array,
+    model: nn.Module,
+    *,
+    max_tokens: int = 256,
+    sampler: Optional[Callable[mx.array, mx.array]] = None,
+    logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None,
+    max_kv_size: Optional[int] = None,
+    prompt_cache: Optional[Any] = None,
+    prefill_step_size: int = 2048,
+    kv_bits: Optional[int] = None,
+    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.
+
+    Args:
+        prompt (mx.array): The input prompt.
+        model (nn.Module): The model to use for generation.
+        max_tokens (int): The maximum number of tokens. Use``-1`` for an infinite
+          generator. Default: ``256``.
+        sampler (Callable[mx.array, mx.array], optional): A sampler for sampling a
+          token from a vector of log probabilities. Default: ``None``.
+        logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
+          A list of functions that take tokens and logits and return the processed
+          logits. Default: ``None``.
+        max_kv_size (int, optional): Maximum size of the key-value cache. Old
+          entries (except the first 4 tokens) will be overwritten.
+        prompt_cache (List[Any], optional): A pre-computed prompt cache. Note, if
+          provided, the cache will be updated in place.
+        prefill_step_size (int): Step size for processing the prompt.
+        kv_bits (int, optional): Number of bits to use for KV cache quantization.
+          None implies no cache quantization. Default: ``None``.
+        kv_group_size (int): Group size for KV cache quantization. Default: ``64``.
+        quantized_kv_start (int): Step to begin using a quantized KV cache.
+           when ``kv_bits`` is non-None. Default: ``0``.
+        prompt_progress_callback (Callable[int, int]): A call-back which takes the
+           prompt tokens processed so far and the total number of prompt tokens.
+        input_embeddings (mx.array, optional): Input embeddings to use in conjunction
+           with prompt tokens. Default: ``None``.
+
+    Yields:
+        Tuple[mx.array, mx.array]: One token and a vector of log probabilities.
+    """
+    if len(prompt) == 0:
+        raise ValueError("Prompt must be non-empty.")
+    if input_embeddings is not None:
+        if not does_model_support_input_embeddings(model):
+            raise ValueError("Model does not support input embeddings.")
+        elif prompt.shape[0] != input_embeddings.shape[0]:
+            raise ValueError(
+                "If using input embeddings, the sequence length must match that of the prompt."
+            )
+
+    tokens = None
+
+    # Create the KV cache for generation
+    if prompt_cache is None:
+        prompt_cache = cache.make_prompt_cache(
+            model,
+            max_kv_size=max_kv_size,
+        )
+
+    prompt_progress_callback = prompt_progress_callback or (lambda *_: None)
+
+    quantize_cache_fn = functools.partial(
+        maybe_quantize_kv_cache,
+        quantized_kv_start=quantized_kv_start,
+        kv_group_size=kv_group_size,
+        kv_bits=kv_bits,
+    )
+
+    sampler = sampler or (lambda x: mx.argmax(x, axis=-1))
+
+    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_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:
+                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)
+            sampled = sampler(logprobs)
+            return sampled, logprobs.squeeze(0)
+
+    with mx.stream(generation_stream):
+        total_prompt_tokens = prompt.shape[0]
+        prompt_processed_tokens = 0
+        while total_prompt_tokens - prompt_processed_tokens > prefill_step_size:
+            _model_call(
+                input_tokens=prompt[:prefill_step_size][None],
+                input_embeddings=(
+                    input_embeddings[:prefill_step_size][None]
+                    if input_embeddings is not None
+                    else None
+                ),
+            )
+            quantize_cache_fn(prompt_cache)
+            mx.eval([c.state for c in prompt_cache])
+            prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens)
+            prompt_processed_tokens += prefill_step_size
+            prompt = prompt[prefill_step_size:]
+            input_embeddings = (
+                input_embeddings[prefill_step_size:]
+                if input_embeddings is not None
+                else input_embeddings
+            )
+            mx.clear_cache()
+
+        y, logprobs = _step(input_tokens=prompt, input_embeddings=input_embeddings)
+
+    mx.async_eval(y, logprobs)
+    n = 0
+    while True:
+        if n != max_tokens:
+            next_y, next_logprobs = _step(y)
+            mx.async_eval(next_y, next_logprobs)
+        if n == 0:
+            mx.eval(y)
+            prompt_progress_callback(total_prompt_tokens, total_prompt_tokens)
+        if n == max_tokens:
+            break
+        yield y.item(), logprobs
+        if n % 256 == 0:
+            mx.clear_cache()
+        y, logprobs = next_y, next_logprobs
+        n += 1
+
+
+def speculative_generate_step(
+    prompt: mx.array,
+    model: nn.Module,
+    draft_model: nn.Module,
+    *,
+    num_draft_tokens=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,
+    prompt_cache: Optional[Any] = None,
+    prefill_step_size: int = 512,
+    kv_bits: Optional[int] = None,
+    kv_group_size: int = 64,
+    quantized_kv_start: int = 0,
+) -> Generator[Tuple[mx.array, mx.array, bool], None, None]:
+    """
+    A generator producing token ids based on the given prompt from the model.
+
+    Args:
+        prompt (mx.array): The input prompt.
+        model (nn.Module): The model to use for generation.
+        draft_model (nn.Module): The draft model for speculative decoding.
+        num_draft_tokens (int, optional): The number of draft tokens for
+          speculative decoding. Default: ``2``.
+        max_tokens (int): The maximum number of tokens. Use``-1`` for an infinite
+          generator. Default: ``256``.
+        sampler (Callable[mx.array, mx.array], optional): A sampler for sampling a
+          token from a vector of log probabilities. Default: ``None``.
+        logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
+          A list of functions that take tokens and logits and return the processed
+          logits. Default: ``None``.
+        prompt_cache (List[Any], optional): A pre-computed prompt cache. Note, if
+          provided, the cache will be updated in place. The cache must be trimmable.
+        prefill_step_size (int): Step size for processing the prompt.
+        kv_bits (int, optional): Number of bits to use for KV cache quantization.
+          None implies no cache quantization. Default: ``None``.
+        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``.
+
+    Yields:
+        Tuple[mx.array, mx.array, bool]: One token, a vector of log probabilities,
+          and a bool indicating if the token was generated by the draft model
+    """
+
+    y = prompt.astype(mx.uint32)
+    prev_tokens = None
+
+    # Create the KV cache for generation
+    if prompt_cache is None:
+        model_cache = cache.make_prompt_cache(model)
+        draft_cache = cache.make_prompt_cache(draft_model)
     else:
-        color = "red"
-    colorprint(color, s)
+        model_cache = prompt_cache[: len(model.layers)]
+        draft_cache = prompt_cache[len(model.layers) :]
+
+    sampler = sampler or (lambda x: mx.argmax(x, axis=-1))
+
+    quantize_cache_fn = functools.partial(
+        maybe_quantize_kv_cache,
+        quantized_kv_start=quantized_kv_start,
+        kv_group_size=kv_group_size,
+        kv_bits=kv_bits,
+    )
+
+    def _process_and_sample(tokens, logits):
+        if logits_processors:
+            for processor in logits_processors:
+                logits = processor(tokens, logits)
+
+        logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
+        y = sampler(logprobs)
+        return y, logprobs
+
+    def _step(model, cache, y, n_predict=1):
+        with mx.stream(generation_stream):
+            logits = model(y[None], cache=cache)
+            logits = logits[:, -n_predict:, :]
+
+            quantize_cache_fn(cache)
+            if logits_processors:
+                nonlocal prev_tokens
+                out_y, out_logprobs = [], []
+                if n_predict > 1:
+                    y = y[: -(n_predict - 1)]
+                for i in range(n_predict):
+                    prev_tokens = (
+                        mx.concat([prev_tokens, y]) if prev_tokens is not None else y
+                    )
+                    y, logprobs = _process_and_sample(prev_tokens, logits[:, i, :])
+                    out_y.append(y)
+                    out_logprobs.append(logprobs)
+                return mx.concatenate(out_y, axis=0), mx.concatenate(
+                    out_logprobs, axis=0
+                )
+            else:
+                return _process_and_sample(None, logits.squeeze(0))
+
+    def _prefill(model, cache, y):
+        while y.size > prefill_step_size:
+            model(y[:prefill_step_size][None], cache=cache)
+            quantize_cache_fn(cache)
+            mx.eval([c.state for c in cache])
+            y = y[prefill_step_size:]
+            mx.clear_cache()
+        return y
+
+    def _rewind_cache(num_draft, num_accept):
+        cache.trim_prompt_cache(model_cache, num_draft - num_accept)
+        cache.trim_prompt_cache(draft_cache, max(num_draft - num_accept - 1, 0))
+
+    def _draft_generate(y, num_draft):
+        if num_draft == 0:
+            return mx.array([], mx.uint32)
+        ys = []
+        for _ in range(num_draft):
+            y, _ = _step(draft_model, draft_cache, y)
+            mx.async_eval(y)
+            ys.append(y)
+        return mx.concatenate(ys)
+
+    with mx.stream(generation_stream):
+        draft_y = _prefill(draft_model, draft_cache, y)
+        y = _prefill(model, model_cache, y)
+
+    ntoks = 0
+    # Set these so the finally block doesn't raise
+    num_draft = 0
+    n = 0
+    try:
+        while True:
+            num_draft = min(max_tokens - ntoks, num_draft_tokens)
+            draft_tokens = _draft_generate(draft_y, num_draft)
+            if prev_tokens is not None:
+                prev_tokens = prev_tokens[: prev_tokens.size - y.size - num_draft + 1]
+            y = mx.concatenate([y, draft_tokens])
+            tokens, logprobs = _step(model, model_cache, y, num_draft + 1)
+            mx.eval(tokens, draft_tokens)
+            draft_tokens = draft_tokens.tolist()
+            tokens = tokens.tolist()
+            n = 0
+            while n < num_draft:
+                tn, dtn, lpn = tokens[n], draft_tokens[n], logprobs[n]
+                if tn != dtn:
+                    break
+                n += 1
+                ntoks += 1
+                yield tn, lpn, True
+                if ntoks == max_tokens:
+                    break
+            if ntoks < max_tokens:
+                ntoks += 1
+                yield tokens[n], logprobs[n], False
+
+            if ntoks == max_tokens:
+                break
+
+            y = mx.array([tokens[n]], mx.uint32)
+            draft_y = y
+
+            # If we accepted all the draft tokens, include the last
+            # draft token in the next draft step since it hasn't been
+            # processed yet by the draft model
+            if n == num_draft:
+                draft_y = mx.concatenate(
+                    [mx.array(draft_tokens[-1:], mx.uint32), draft_y]
+                )
+
+            if prev_tokens is not None:
+                prev_tokens = prev_tokens[: -max(num_draft - n, 1)]
+            _rewind_cache(num_draft, n)
+    finally:
+        _rewind_cache(num_draft, n)
+
+
+def stream_generate(
+    model: nn.Module,
+    tokenizer: Union[PreTrainedTokenizer, TokenizerWrapper],
+    prompt: Union[str, mx.array, List[int]],
+    draft_model: Optional[nn.Module] = None,
+    **kwargs,
+) -> Generator[GenerationResponse, None, None]:
+    """
+    A generator producing text based on the given prompt from the model.
+
+    Args:
+        model (nn.Module): The model to use for generation.
+        tokenizer (PreTrainedTokenizer): The tokenizer.
+        prompt (Union[str, mx.array, List[int]]): The input prompt string or
+          integer tokens.
+        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``.
+        kwargs: The remaining options get passed to :func:`generate_step`.
+          See :func:`generate_step` for more details.
+
+    Yields:
+        GenerationResponse: An instance containing the generated text segment and
+            associated metadata. See :class:`GenerationResponse` for details.
+    """
+    if not isinstance(tokenizer, TokenizerWrapper):
+        tokenizer = TokenizerWrapper(tokenizer)
+
+    if not isinstance(prompt, mx.array):
+        if isinstance(prompt, str):
+            # Try to infer if special tokens are needed
+            add_special_tokens = tokenizer.bos_token is None or not prompt.startswith(
+                tokenizer.bos_token
+            )
+            prompt = tokenizer.encode(prompt, add_special_tokens=add_special_tokens)
+        prompt = mx.array(prompt)
+
+    detokenizer = tokenizer.detokenizer
+
+    if draft_model is None:
+        kwargs.pop("num_draft_tokens", None)
+        token_generator = generate_step(prompt, model, **kwargs)
+        # from_draft always false for non-speculative generation
+        token_generator = (
+            (token, logprobs, False) for token, logprobs in token_generator
+        )
+    else:
+        kwargs.pop("max_kv_size", 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:
+                prompt_time = time.perf_counter() - tic
+                prompt_tps = prompt.size / prompt_time
+                tic = time.perf_counter()
+            if token in tokenizer.eos_token_ids:
+                break
+
+            detokenizer.add_token(token)
+
+            yield GenerationResponse(
+                text=detokenizer.last_segment,
+                token=token,
+                logprobs=logprobs,
+                from_draft=from_draft,
+                prompt_tokens=prompt.size,
+                prompt_tps=prompt_tps,
+                generation_tokens=n + 1,
+                generation_tps=(n + 1) / (time.perf_counter() - tic),
+                peak_memory=mx.get_peak_memory() / 1e9,
+                finish_reason=None,
+            )
+
+        detokenizer.finalize()
+        yield GenerationResponse(
+            text=detokenizer.last_segment,
+            token=token,
+            logprobs=logprobs,
+            from_draft=from_draft,
+            prompt_tokens=prompt.size,
+            prompt_tps=prompt_tps,
+            generation_tokens=n + 1,
+            generation_tps=(n + 1) / (time.perf_counter() - tic),
+            peak_memory=mx.get_peak_memory() / 1e9,
+            finish_reason="stop" if token in tokenizer.eos_token_ids else "length",
+        )
+
+
+def generate(
+    model: nn.Module,
+    tokenizer: Union[PreTrainedTokenizer, TokenizerWrapper],
+    prompt: Union[str, List[int]],
+    verbose: bool = False,
+    formatter: Optional[Callable] = None,
+    **kwargs,
+) -> str:
+    """
+    Generate a complete response from the model.
+
+    Args:
+       model (nn.Module): The language model.
+       tokenizer (PreTrainedTokenizer): The tokenizer.
+       prompt (Union[str, List[int]]): The input prompt string or integer tokens.
+       verbose (bool): If ``True``, print tokens and timing information.
+           Default: ``False``.
+       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)
+
+    text = ""
+    for response in stream_generate(model, tokenizer, prompt, **kwargs):
+        if verbose:
+            print(response.text, end="", flush=True)
+        text += response.text
+
+    if verbose:
+        print()
+        print("=" * 10)
+        if len(text) == 0:
+            print("No text generated for this prompt")
+            return
+        print(
+            f"Prompt: {response.prompt_tokens} tokens, "
+            f"{response.prompt_tps:.3f} tokens-per-sec"
+        )
+        print(
+            f"Generation: {response.generation_tokens} tokens, "
+            f"{response.generation_tps:.3f} tokens-per-sec"
+        )
+        print(f"Peak memory: {response.peak_memory:.3f} GB")
+    return text
 
 
 def main():
     parser = setup_arg_parser()
     args = parser.parse_args()
 
-    mx.random.seed(args.seed)
+    if args.seed is not None:
+        mx.random.seed(args.seed)
 
-    if args.cache_limit_gb is not None:
-        mx.metal.set_cache_limit(args.cache_limit_gb * 1024 * 1024 * 1024)
+    # Load the prompt cache and metadata if a cache file is provided
+    using_cache = args.prompt_cache_file is not None
+    if using_cache:
+        prompt_cache, metadata = load_prompt_cache(
+            args.prompt_cache_file,
+            return_metadata=True,
+        )
+        if isinstance(prompt_cache[0], QuantizedKVCache):
+            if args.kv_bits is not None and args.kv_bits != prompt_cache[0].bits:
+                raise ValueError(
+                    "--kv-bits does not match the kv cache loaded from --prompt-cache-file."
+                )
+            if args.kv_group_size != prompt_cache[0].group_size:
+                raise ValueError(
+                    "--kv-group-size does not match the kv cache loaded from --prompt-cache-file."
+                )
 
     # Building tokenizer_config
-    tokenizer_config = {"trust_remote_code": True if args.trust_remote_code else None}
-    if args.eos_token is not None:
-        tokenizer_config["eos_token"] = args.eos_token
+    tokenizer_config = (
+        {} if not using_cache else json.loads(metadata["tokenizer_config"])
+    )
+    tokenizer_config["trust_remote_code"] = True
+
+    model_path = args.model
+    if using_cache:
+        if model_path is None:
+            model_path = metadata["model"]
+        elif model_path != metadata["model"]:
+            raise ValueError(
+                f"Providing a different model ({model_path}) than that "
+                f"used to create the prompt cache ({metadata['model']}) "
+                "is an error."
+            )
+    model_path = model_path or DEFAULT_MODEL
 
     model, tokenizer = load(
-        args.model,
+        model_path,
         adapter_path=args.adapter_path,
         tokenizer_config=tokenizer_config,
     )
+    for eos_token in args.extra_eos_token:
+        tokenizer.add_eos_token(eos_token)
+
+    template_kwargs = {}
+    if args.chat_template_config is not None:
+        template_kwargs = json.loads(args.chat_template_config)
 
     if args.use_default_chat_template:
         if tokenizer.chat_template is None:
             tokenizer.chat_template = tokenizer.default_chat_template
+    elif using_cache:
+        tokenizer.chat_template = json.loads(metadata["chat_template"])
 
-    if not args.ignore_chat_template and (
-        hasattr(tokenizer, "apply_chat_template")
-        and tokenizer.chat_template is not None
-    ):
-        messages = [{"role": "user", "content": args.prompt}]
+    prompt = args.prompt.replace("\\n", "\n").replace("\\t", "\t")
+    prompt = sys.stdin.read() if prompt == "-" else prompt
+    if not args.ignore_chat_template and tokenizer.chat_template is not None:
+        if args.system_prompt is not None:
+            messages = [{"role": "system", "content": args.system_prompt}]
+        else:
+            messages = []
+        messages.append({"role": "user", "content": prompt})
+
+        has_prefill = args.prefill_response is not None
+        if has_prefill:
+            messages.append({"role": "assistant", "content": args.prefill_response})
         prompt = tokenizer.apply_chat_template(
-            messages, tokenize=False, add_generation_prompt=True
+            messages,
+            tokenize=False,
+            continue_final_message=has_prefill,
+            add_generation_prompt=not has_prefill,
+            **template_kwargs,
         )
+
+        # Treat the prompt as a suffix assuming that the prefix is in the
+        # stored kv cache.
+        if using_cache:
+            messages[-1]["content"] = "<query>"
+            test_prompt = tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                continue_final_message=has_prefill,
+                add_generation_prompt=not has_prefill,
+            )
+            prompt = prompt[test_prompt.index("<query>") :]
+        prompt = tokenizer.encode(prompt, add_special_tokens=False)
     else:
-        prompt = args.prompt
+        prompt = tokenizer.encode(prompt)
 
-    formatter = colorprint_by_t0 if args.colorize else None
-
-    generate(
+    if args.draft_model is not None:
+        draft_model, draft_tokenizer = load(args.draft_model)
+        if draft_tokenizer.vocab_size != tokenizer.vocab_size:
+            raise ValueError("Draft model tokenizer does not match model tokenizer.")
+    else:
+        draft_model = None
+    sampler = make_sampler(
+        args.temp,
+        args.top_p,
+        args.min_p,
+        args.min_tokens_to_keep,
+        top_k=args.top_k,
+        xtc_probability=args.xtc_probability,
+        xtc_threshold=args.xtc_threshold,
+        xtc_special_tokens=tokenizer.encode("\n") + list(tokenizer.eos_token_ids),
+    )
+    response = generate(
         model,
         tokenizer,
         prompt,
-        args.max_tokens,
-        verbose=True,
-        formatter=formatter,
-        temp=args.temp,
-        top_p=args.top_p,
+        max_tokens=args.max_tokens,
+        verbose=args.verbose,
+        sampler=sampler,
+        max_kv_size=args.max_kv_size,
+        prompt_cache=prompt_cache if using_cache else None,
+        kv_bits=args.kv_bits,
+        kv_group_size=args.kv_group_size,
+        quantized_kv_start=args.quantized_kv_start,
+        draft_model=draft_model,
+        num_draft_tokens=args.num_draft_tokens,
     )
+    if not args.verbose:
+        print(response)
 
 
 if __name__ == "__main__":
+    print(
+        "Calling `python -m mlx_lm.generate...` directly is deprecated."
+        " Use `mlx_lm.generate...` or `python -m mlx_lm generate ...` instead."
+    )
     main()

mlx_lm/gguf.py

@@ -59,7 +59,7 @@ class HfVocab:
         for token_id in range(self.vocab_size_base):
             if token_id in self.added_tokens_ids:
                 continue
-            token_text = reverse_vocab[token_id].encode("utf-8")
+            token_text = reverse_vocab[token_id]
             yield token_text, self.get_token_score(token_id), self.get_token_type(
                 token_id, token_text, self.special_ids
             )
@@ -67,7 +67,7 @@ class HfVocab:
     def get_token_type(
         self, token_id: int, token_text: bytes, special_ids: Set[int]
     ) -> TokenType:
-        if re.fullmatch(rb"<0x[0-9A-Fa-f]{2}>", token_text):
+        if re.fullmatch(r"<0x[0-9A-Fa-f]{2}>", token_text):
             return TokenType.BYTE
         return TokenType.CONTROL if token_id in special_ids else TokenType.NORMAL
 
@@ -77,14 +77,12 @@ class HfVocab:
     def added_tokens(self) -> Iterable[Tuple[bytes, float, TokenType]]:
         for text in self.added_tokens_list:
             if text in self.specials:
-                toktype = self.get_token_type(
-                    self.specials[text], b"", self.special_ids
-                )
+                toktype = self.get_token_type(self.specials[text], "", self.special_ids)
                 score = self.get_token_score(self.specials[text])
             else:
                 toktype = TokenType.USER_DEFINED
                 score = -1000.0
-            yield text.encode("utf-8"), score, toktype
+            yield text, score, toktype
 
     def has_newline_token(self):
         return "<0x0A>" in self.tokenizer.vocab or "\n" in self.tokenizer.vocab
@@ -243,15 +241,18 @@ def prepare_metadata(config, vocab):
     metadata["tokenizer.ggml.tokens"] = tokens
     metadata["tokenizer.ggml.scores"] = mx.array(scores, dtype=mx.float32)
     metadata["tokenizer.ggml.token_type"] = mx.array(toktypes, dtype=mx.uint32)
-    metadata["tokenizer.ggml.bos_token_id"] = mx.array(
-        vocab.tokenizer.bos_token_id, dtype=mx.uint32
-    )
-    metadata["tokenizer.ggml.eos_token_id"] = mx.array(
-        vocab.tokenizer.eos_token_id, dtype=mx.uint32
-    )
-    metadata["tokenizer.ggml.unknown_token_id"] = mx.array(
-        vocab.tokenizer.unk_token_id, dtype=mx.uint32
-    )
+    if vocab.tokenizer.bos_token_id is not None:
+        metadata["tokenizer.ggml.bos_token_id"] = mx.array(
+            vocab.tokenizer.bos_token_id, dtype=mx.uint32
+        )
+    if vocab.tokenizer.eos_token_id is not None:
+        metadata["tokenizer.ggml.eos_token_id"] = mx.array(
+            vocab.tokenizer.eos_token_id, dtype=mx.uint32
+        )
+    if vocab.tokenizer.unk_token_id is not None:
+        metadata["tokenizer.ggml.unknown_token_id"] = mx.array(
+            vocab.tokenizer.unk_token_id, dtype=mx.uint32
+        )
 
     metadata = {k: v for k, v in metadata.items() if v is not None}
     return metadata

mlx_lm/lora.py

@@ -1,23 +1,23 @@
-# Copyright © 2024 Apple Inc.
-
 import argparse
 import math
+import os
 import re
 import types
 from pathlib import Path
 
+import mlx.core as mx
 import mlx.nn as nn
 import mlx.optimizers as optim
 import numpy as np
 import yaml
 
-from .tokenizer_utils import TokenizerWrapper
-from .tuner.datasets import load_dataset
+from .tuner.callbacks import WandBCallback
+from .tuner.datasets import CacheDataset, load_dataset
 from .tuner.trainer import TrainingArgs, TrainingCallback, evaluate, train
 from .tuner.utils import (
-    apply_lora_layers,
     build_schedule,
     linear_to_lora_layers,
+    load_adapters,
     print_trainable_parameters,
 )
 from .utils import load, save_config
@@ -41,9 +41,15 @@ yaml_loader.add_implicit_resolver(
 CONFIG_DEFAULTS = {
     "model": "mlx_model",
     "train": False,
+    "fine_tune_type": "lora",
+    "optimizer": "adam",
+    "optimizer_config": {
+        "adam": {},
+        "adamw": {},
+    },
     "data": "data/",
     "seed": 0,
-    "lora_layers": 16,
+    "num_layers": 16,
     "batch_size": 4,
     "iters": 1000,
     "val_batches": 25,
@@ -56,9 +62,12 @@ CONFIG_DEFAULTS = {
     "test": False,
     "test_batches": 500,
     "max_seq_length": 2048,
+    "config": None,
+    "grad_checkpoint": False,
     "lr_schedule": None,
-    "lora_parameters": {"rank": 8, "alpha": 16, "dropout": 0.0, "scale": 10.0},
-    "use_dora": False,
+    "lora_parameters": {"rank": 8, "dropout": 0.0, "scale": 20.0},
+    "mask_prompt": False,
+    "wandb": None,
 }
 
 
@@ -66,6 +75,7 @@ def build_parser():
     parser = argparse.ArgumentParser(description="LoRA or QLoRA finetuning.")
     parser.add_argument(
         "--model",
+        type=str,
         help="The path to the local model directory or Hugging Face repo.",
     )
 
@@ -79,10 +89,32 @@ def build_parser():
     parser.add_argument(
         "--data",
         type=str,
-        help="Directory with {train, valid, test}.jsonl files",
+        help=(
+            "Directory with {train, valid, test}.jsonl files or the name "
+            "of a Hugging Face dataset (e.g., 'mlx-community/wikisql')"
+        ),
     )
     parser.add_argument(
-        "--lora-layers",
+        "--fine-tune-type",
+        type=str,
+        choices=["lora", "dora", "full"],
+        help="Type of fine-tuning to perform: lora, dora, or full.",
+    )
+    parser.add_argument(
+        "--optimizer",
+        type=str,
+        choices=["adam", "adamw"],
+        default=None,
+        help="Optimizer to use for training: adam or adamw",
+    )
+    parser.add_argument(
+        "--mask-prompt",
+        action="store_true",
+        help="Mask the prompt in the loss when training",
+        default=None,
+    )
+    parser.add_argument(
+        "--num-layers",
         type=int,
         help="Number of layers to fine-tune. Default is 16, use -1 for all.",
     )
@@ -107,12 +139,12 @@ def build_parser():
     parser.add_argument(
         "--resume-adapter-file",
         type=str,
-        help="Load path to resume training with the given adapters.",
+        help="Load path to resume training from the given fine-tuned weights.",
     )
     parser.add_argument(
         "--adapter-path",
         type=str,
-        help="Save/load path for the adapters.",
+        help="Save/load path for the fine-tuned weights.",
     )
     parser.add_argument(
         "--save-every",
@@ -138,7 +170,7 @@ def build_parser():
     parser.add_argument(
         "-c",
         "--config",
-        default=None,
+        type=str,
         help="A YAML configuration file with the training options",
     )
     parser.add_argument(
@@ -147,36 +179,57 @@ def build_parser():
         help="Use gradient checkpointing to reduce memory use.",
         default=None,
     )
-    parser.add_argument("--seed", type=int, default=None, help="The PRNG seed")
     parser.add_argument(
-        "--use-dora", action="store_true", default=None, help="Use DoRA to finetune."
+        "--wandb",
+        type=str,
+        default=None,
+        help="WandB project name to report training metrics. Disabled if None.",
     )
+    parser.add_argument("--seed", type=int, help="The PRNG seed")
     return parser
 
 
 def train_model(
     args,
     model: nn.Module,
-    tokenizer: TokenizerWrapper,
     train_set,
     valid_set,
     training_callback: TrainingCallback = None,
 ):
-    # Freeze all layers
+    mx.random.seed(args.seed)
     model.freeze()
+    if args.num_layers > len(model.layers):
+        raise ValueError(
+            f"Requested to train {args.num_layers} layers "
+            f"but the model only has {len(model.layers)} layers."
+        )
 
-    # Convert linear layers to lora layers and unfreeze in the process
-    linear_to_lora_layers(model, args.lora_layers, args.lora_parameters)
+    if args.fine_tune_type == "full":
+        for l in model.layers[-max(args.num_layers, 0) :]:
+            l.unfreeze()
 
-    # Resume training the given adapters.
+        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(
+            model,
+            args.num_layers,
+            args.lora_parameters,
+            use_dora=(args.fine_tune_type == "dora"),
+        )
+    else:
+        raise ValueError(f"Received unknown fine-tune-type {args.fine_tune_type}")
+
+    # Resume from weights if provided
     if args.resume_adapter_file is not None:
-        print(f"Loading pretrained adapters from {args.resume_adapter_file}")
+        print(f"Loading fine-tuned weights from {args.resume_adapter_file}")
         model.load_weights(args.resume_adapter_file, strict=False)
 
     print_trainable_parameters(model)
 
     adapter_path = Path(args.adapter_path)
     adapter_path.mkdir(parents=True, exist_ok=True)
+
     adapter_file = adapter_path / "adapters.safetensors"
     save_config(vars(args), adapter_path / "adapter_config.json")
 
@@ -193,31 +246,36 @@ def train_model(
         grad_checkpoint=args.grad_checkpoint,
     )
 
-    model.train()
-    opt = optim.Adam(
-        learning_rate=(
-            build_schedule(args.lr_schedule) if args.lr_schedule else args.learning_rate
-        )
-    )
+    # Initialize the selected optimizer
+    lr = build_schedule(args.lr_schedule) if args.lr_schedule else args.learning_rate
+
+    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
+    else:
+        raise ValueError(f"Unsupported optimizer: {optimizer_name}")
+
+    opt = opt_class(learning_rate=lr, **optimizer_config)
+
     # Train model
     train(
         model=model,
-        tokenizer=tokenizer,
         args=training_args,
         optimizer=opt,
-        train_dataset=train_set,
-        val_dataset=valid_set,
+        train_dataset=CacheDataset(train_set),
+        val_dataset=CacheDataset(valid_set),
         training_callback=training_callback,
     )
 
 
-def evaluate_model(args, model: nn.Module, tokenizer: TokenizerWrapper, test_set):
-    model.eval()
-
+def evaluate_model(args, model: nn.Module, test_set):
     test_loss = evaluate(
         model=model,
-        dataset=test_set,
-        tokenizer=tokenizer,
+        dataset=CacheDataset(test_set),
         batch_size=args.batch_size,
         num_batches=args.test_batches,
         max_seq_length=args.max_seq_length,
@@ -231,6 +289,14 @@ def evaluate_model(args, model: nn.Module, tokenizer: TokenizerWrapper, test_set
 def run(args, training_callback: TrainingCallback = None):
     np.random.seed(args.seed)
 
+    if args.wandb is not None:
+        training_callback = WandBCallback(
+            project_name=args.wandb,
+            log_dir=args.adapter_path,
+            config=vars(args),
+            wrapped_callback=training_callback,
+        )
+
     print("Loading pretrained model")
     model, tokenizer = load(args.model)
 
@@ -240,20 +306,21 @@ def run(args, training_callback: TrainingCallback = None):
     if args.test and not args.train:
         # Allow testing without LoRA layers by providing empty path
         if args.adapter_path != "":
-            apply_lora_layers(model, args.adapter_path)
+            load_adapters(model, args.adapter_path)
 
     elif args.train:
         print("Training")
-        train_model(args, model, tokenizer, train_set, valid_set, training_callback)
+        train_model(args, model, train_set, valid_set, training_callback)
     else:
         raise ValueError("Must provide at least one of --train or --test")
 
     if args.test:
         print("Testing")
-        evaluate_model(args, model, tokenizer, test_set)
+        evaluate_model(args, model, test_set)
 
 
 def main():
+    os.environ["TOKENIZERS_PARALLELISM"] = "true"
     parser = build_parser()
     args = parser.parse_args()
     config = args.config
@@ -275,4 +342,8 @@ def main():
 
 
 if __name__ == "__main__":
+    print(
+        "Calling `python -m mlx_lm.lora...` directly is deprecated."
+        " Use `mlx_lm.lora...` or `python -m mlx_lm lora ...` instead."
+    )
     main()

mlx_lm/manage.py

@@ -2,23 +2,37 @@ import argparse
 from typing import List, Union
 
 from huggingface_hub import scan_cache_dir
-from transformers.commands.user import tabulate
+
+
+def tabulate(rows: List[List[Union[str, int]]], headers: List[str]) -> str:
+    """
+    Inspired by:
+    - stackoverflow.com/a/8356620/593036
+    - stackoverflow.com/questions/9535954/printing-lists-as-tabular-data
+    """
+    col_widths = [max(len(str(x)) for x in col) for col in zip(*rows, headers)]
+    row_format = ("{{:{}}} " * len(headers)).format(*col_widths)
+    lines = []
+    lines.append(row_format.format(*headers))
+    lines.append(row_format.format(*["-" * w for w in col_widths]))
+    for row in rows:
+        lines.append(row_format.format(*row))
+    return "\n".join(lines)
 
 
 def ask_for_confirmation(message: str) -> bool:
+    """Ask user for confirmation with Y/N prompt.
+    Returns True for Y/yes, False for N/no/empty."""
     y = ("y", "yes", "1")
-    n = ("n", "no", "0")
-    all_values = y + n + ("",)
-    full_message = f"{message} (Y/n) "
+    n = ("n", "no", "0", "")
+    full_message = f"{message} (y/n) "
     while True:
         answer = input(full_message).lower()
-        if answer == "":
-            return False
         if answer in y:
             return True
         if answer in n:
             return False
-        print(f"Invalid input. Must be one of {all_values}")
+        print(f"Invalid input. Must be one of: yes/no/y/n or empty for no")
 
 
 def main():
@@ -43,9 +57,7 @@ def main():
     args = parser.parse_args()
 
     if args.scan:
-        print(
-            "Scanning Hugging Face cache for models with" f'pattern "{args.pattern}".'
-        )
+        print(f'Scanning Hugging Face cache for models with pattern "{args.pattern}".')
         hf_cache_info = scan_cache_dir()
         print(
             tabulate(
@@ -86,36 +98,46 @@ def main():
             if args.pattern in repo.repo_id
         ]
         if repos:
+            print("\nFound the following models:")
             print(
                 tabulate(
                     rows=[
                         [
                             repo.repo_id,
+                            repo.size_on_disk_str,  # Added size information
                             str(repo.repo_path),
                         ]
                         for repo in repos
                     ],
                     headers=[
                         "REPO ID",
+                        "SIZE",  # Added size header
                         "LOCAL PATH",
                     ],
                 )
             )
 
-            confirmed = ask_for_confirmation(f"Confirm deletion ?")
+            confirmed = ask_for_confirmation(
+                "\nAre you sure you want to delete these models?"
+            )
             if confirmed:
                 for model_info in repos:
+                    print(f"\nDeleting {model_info.repo_id}...")
                     for revision in sorted(
                         model_info.revisions, key=lambda revision: revision.commit_hash
                     ):
                         strategy = hf_cache_info.delete_revisions(revision.commit_hash)
                         strategy.execute()
-                print("Model(s) deleted.")
+                print("\nModel(s) deleted successfully.")
             else:
-                print("Deletion is cancelled. Do nothing.")
+                print("\nDeletion cancelled - no changes made.")
         else:
-            print(f"No models found.")
+            print(f'No models found matching pattern "{args.pattern}"')
 
 
 if __name__ == "__main__":
+    print(
+        "Calling `python -m mlx_lm.manage...` directly is deprecated."
+        " Use `mlx_lm.manage...` or `python -m mlx_lm manage ...` instead."
+    )
     main()

mlx_lm/merge.py

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

mlx_lm/models/afm7.py

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

mlx_lm/models/baichuan_m1.py

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

mlx_lm/models/base.py

@@ -1,46 +1,13 @@
+# Copyright © 2023-2024 Apple Inc.
+
 import inspect
 from dataclasses import dataclass
+from typing import Any, Optional
 
 import mlx.core as mx
+from mlx.utils import tree_map
 
-
-def create_additive_causal_mask(N: int, offset: int = 0):
-    rinds = mx.arange(offset + N)
-    linds = mx.arange(offset, offset + N) if offset else rinds
-    mask = linds[:, None] < rinds[None]
-    return mask * -1e9
-
-
-class KVCache:
-
-    def __init__(self, head_dim, n_kv_heads):
-        self.n_kv_heads = n_kv_heads
-        self.head_dim = head_dim
-        self.keys = None
-        self.values = None
-        self.offset = 0
-        self.step = 256
-
-    def update_and_fetch(self, keys, values):
-        prev = self.offset
-        if self.keys is None or (prev + keys.shape[2]) > self.keys.shape[2]:
-            n_steps = (self.step + keys.shape[2] - 1) // self.step
-            shape = (1, self.n_kv_heads, n_steps * self.step, self.head_dim)
-            new_k = mx.zeros(shape, keys.dtype)
-            new_v = mx.zeros(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.keys[..., prev : self.offset, :] = keys
-        self.values[..., prev : self.offset, :] = values
-        return self.keys[..., : self.offset, :], self.values[..., : self.offset, :]
+from .cache import QuantizedKVCache
 
 
 @dataclass
@@ -54,3 +21,113 @@ class BaseModelArgs:
                 if k in inspect.signature(cls).parameters
             }
         )
+
+
+def create_causal_mask(
+    N: int,
+    offset: int = 0,
+    window_size: Optional[int] = None,
+    lengths: Optional[mx.array] = None,
+):
+    rinds = mx.arange(offset + N)
+    linds = mx.arange(offset, offset + N) if offset else rinds
+    linds = linds[:, None]
+    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)
+    return mask
+
+
+def create_attention_mask(
+    h: mx.array, cache: Optional[Any] = 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
+
+
+def quantized_scaled_dot_product_attention(
+    queries: mx.array,
+    q_keys: tuple[mx.array, mx.array, mx.array],
+    q_values: tuple[mx.array, mx.array, mx.array],
+    scale: float,
+    mask: Optional[mx.array],
+    group_size: int = 64,
+    bits: int = 8,
+) -> mx.array:
+    B, n_q_heads, L, D = queries.shape
+    n_kv_heads = q_keys[0].shape[-3]
+    n_repeats = n_q_heads // n_kv_heads
+
+    queries *= scale
+
+    if n_repeats > 1:
+        queries = mx.reshape(queries, (B, n_kv_heads, n_repeats, L, D))
+        q_keys = tree_map(lambda x: mx.expand_dims(x, axis=-3), q_keys)
+        q_values = tree_map(lambda x: mx.expand_dims(x, axis=-3), q_values)
+
+    scores = mx.quantized_matmul(
+        queries, *q_keys, transpose=True, group_size=group_size, bits=bits
+    )
+    if mask is not None:
+        if isinstance(mask, str):
+            qL, kL = scores.shape[-2:]
+            q_indices = mx.arange(kL - qL, kL)
+            k_indices = mx.arange(kL)
+            mask = q_indices[:, None] >= k_indices[None]
+        if mask.dtype == mx.bool_:
+            scores = mx.where(mask, scores, mx.finfo(scores.dtype).min)
+        else:
+            scores += mask
+    scores = mx.softmax(scores, axis=-1, precise=True)
+    out = mx.quantized_matmul(
+        scores, *q_values, transpose=False, group_size=group_size, bits=bits
+    )
+
+    if n_repeats > 1:
+        out = mx.reshape(out, (B, n_q_heads, L, D))
+
+    return out
+
+
+def scaled_dot_product_attention(
+    queries,
+    keys,
+    values,
+    cache,
+    scale: float,
+    mask: Optional[mx.array],
+) -> mx.array:
+    if isinstance(cache, QuantizedKVCache):
+        return quantized_scaled_dot_product_attention(
+            queries,
+            keys,
+            values,
+            scale=scale,
+            mask=mask,
+            group_size=cache.group_size,
+            bits=cache.bits,
+        )
+    else:
+        return mx.fast.scaled_dot_product_attention(
+            queries, keys, values, scale=scale, mask=mask
+        )

mlx_lm/models/bitlinear_layers.py

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

mlx_lm/models/bitnet.py

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

mlx_lm/models/cache.py

@@ -0,0 +1,545 @@
+# Copyright © 2023-2024 Apple Inc.
+
+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
+
+
+def make_prompt_cache(
+    model: nn.Module,
+    max_kv_size: Optional[int] = None,
+) -> List[Any]:
+    """
+    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.
+
+    Args:
+        model (nn.Module): The language model.
+        max_kv_size (Optional[int]): If provided and the model does not have a
+            ``make_cache`` method, a ``RotatingKVCache`` is used with a maximum
+            size of ``max_kv_size``
+    """
+    if hasattr(model, "make_cache"):
+        return model.make_cache()
+
+    num_layers = len(model.layers)
+    if max_kv_size is not None:
+        return [
+            RotatingKVCache(max_size=max_kv_size, keep=4) for _ in range(num_layers)
+        ]
+    else:
+        return [KVCache() for _ in range(num_layers)]
+
+
+def save_prompt_cache(file_name: str, cache: List[Any], metadata: Dict[str, str] = {}):
+    """
+    Save a pre-computed prompt cache to a file.
+
+    Args:
+        file_name (str): The ``.safetensors`` file name.
+        cache (List[Any]): The model state.
+        metadata (Dict[str, str]): Optional metadata to save along with model
+            state.
+    """
+    cache_data = [c.state for c in cache]
+    cache_info = [c.meta_state for c in cache]
+    cache_data = dict(tree_flatten(cache_data))
+    cache_classes = [type(c).__name__ for c in cache]
+    cache_metadata = [cache_info, metadata, cache_classes]
+    cache_metadata = dict(tree_flatten(cache_metadata))
+    mx.save_safetensors(file_name, cache_data, cache_metadata)
+
+
+def load_prompt_cache(file_name, return_metadata=False):
+    """
+    Load a prompt cache from a file.
+
+    Args:
+        file_name (str): The ``.safetensors`` file name.
+        return_metadata (bool): Whether or not to return metadata.
+            Default: ``False``.
+
+    Returns:
+        List[Any] or Tuple[List[Any], Dict[str, str]]: The prompt cache and
+            the metadata if requested.
+    """
+    arrays, cache_metadata = mx.load(file_name, return_metadata=True)
+    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
+    if return_metadata:
+        return cache, metadata
+    return cache
+
+
+def can_trim_prompt_cache(cache: List[Any]) -> bool:
+    """
+    Check if model's cache can be trimmed.
+    """
+    return all(c.is_trimmable() for c in cache)
+
+
+def trim_prompt_cache(cache: List[Any], num_tokens: int) -> List[Any]:
+    """
+    Trim the model's cache by the given number of tokens.
+
+    This function will trim the cache if possible (in-place) and return the
+    number of tokens that were trimmed.
+
+    Args:
+        cache (List[Any]): The model's cache.
+        num_tokens (int): The number of tokens to trim.
+
+    Returns:
+        (int): The number of tokens that were trimmed.
+    """
+    if not can_trim_prompt_cache(cache) or len(cache) == 0:
+        return 0
+    return [c.trim(num_tokens) for c in cache][0]
+
+
+class _BaseCache:
+    @property
+    def state(self):
+        return []
+
+    @state.setter
+    def state(self, v):
+        if v is not None and v:
+            raise ValueError("This cache has no state but a state was set.")
+
+    @property
+    def meta_state(self):
+        return ""
+
+    @meta_state.setter
+    def meta_state(self, v):
+        if v is not None and v:
+            raise ValueError("This cache has no meta_state but a meta_state was set.")
+
+    def is_trimmable(self):
+        return False
+
+
+class ConcatenateKVCache(_BaseCache):
+    """ConcatenateKVCache the simplest KV cache implementation.
+
+    Can be used as a mock KV cache or when large blocks are being processed at
+    a time in which case KVCache isn't necessarily faster. Consider using the
+    KVCache with a larger step size before using this cache.
+    """
+
+    def __init__(self):
+        self.keys = None
+        self.values = None
+        self.offset = 0
+
+    def update_and_fetch(self, keys, values):
+        if self.keys is None:
+            self.keys = keys
+            self.values = values
+        else:
+            self.keys = mx.concatenate([self.keys, keys], axis=-2)
+            self.values = mx.concatenate([self.values, values], axis=-2)
+        self.offset = self.keys.shape[-2]
+
+        return self.keys, self.values
+
+    @property
+    def state(self):
+        return self.keys, self.values
+
+    @state.setter
+    def state(self, v):
+        self.keys, self.values = v
+        self.offset = self.keys.shape[-2]
+
+
+class QuantizedKVCache(_BaseCache):
+    def __init__(self, group_size: int = 64, bits: int = 8):
+        self.keys = None
+        self.values = None
+        self.offset = 0
+        self.step = 256
+        self.group_size = group_size
+        self.bits = bits
+
+    def update_and_fetch(self, keys, values):
+        B, n_kv_heads, num_steps, k_head_dim = keys.shape
+        v_head_dim = values.shape[-1]
+        prev = self.offset
+
+        if self.keys is None or (prev + num_steps) > self.keys[0].shape[-2]:
+            el_per_int = 8 * mx.uint32.size // self.bits
+            new_steps = (self.step + num_steps - 1) // self.step * self.step
+            shape = (B, n_kv_heads, new_steps)
+
+            def init_quant(dim):
+                return (
+                    mx.zeros((*shape, dim // el_per_int), dtype=mx.uint32),
+                    mx.zeros((*shape, dim // self.group_size), dtype=keys.dtype),
+                    mx.zeros((*shape, dim // self.group_size), dtype=keys.dtype),
+                )
+
+            def expand_quant(x):
+                new_x = mx.zeros((*shape, x.shape[-1]), dtype=x.dtype)
+                return mx.concatenate([x, new_x], axis=-2)
+
+            if self.keys is not None:
+                if prev % self.step != 0:
+                    self.keys, self.values = tree_map(
+                        lambda x: x[..., :prev, :], (self.keys, self.values)
+                    )
+
+                self.keys, self.values = tree_map(
+                    expand_quant, (self.keys, self.values)
+                )
+            else:
+                self.keys, self.values = init_quant(k_head_dim), init_quant(v_head_dim)
+
+        self.offset += num_steps
+
+        keys = mx.quantize(keys, group_size=self.group_size, bits=self.bits)
+        values = mx.quantize(values, group_size=self.group_size, bits=self.bits)
+        for i in range(len(self.keys)):
+            self.keys[i][..., prev : self.offset, :] = keys[i]
+            self.values[i][..., prev : self.offset, :] = values[i]
+
+        return tree_map(lambda x: x[..., : self.offset, :], (self.keys, self.values))
+
+    @property
+    def state(self):
+        if self.offset == self.keys[0].shape[2]:
+            return self.keys, self.values
+        else:
+            return tree_map(
+                lambda x: x[..., : self.offset, :], (self.keys, self.values)
+            )
+
+    @state.setter
+    def state(self, v):
+        self.keys, self.values = v
+
+    @property
+    def meta_state(self):
+        return tuple(map(str, (self.step, 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)
+
+    def is_trimmable(self):
+        return True
+
+    def trim(self, n):
+        n = min(self.offset, n)
+        self.offset -= n
+        return n
+
+
+class KVCache(_BaseCache):
+    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
+        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.keys[..., prev : self.offset, :] = keys
+        self.values[..., prev : self.offset, :] = values
+        return self.keys[..., : self.offset, :], self.values[..., : self.offset, :]
+
+    @property
+    def state(self):
+        if self.offset == self.keys.shape[2]:
+            return self.keys, self.values
+        else:
+            return (
+                self.keys[..., : self.offset, :],
+                self.values[..., : self.offset, :],
+            )
+
+    @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 to_quantized(self, group_size: int = 64, bits: int = 4) -> QuantizedKVCache:
+        quant_cache = QuantizedKVCache(group_size=group_size, bits=bits)
+        quant_cache.offset = self.offset
+        if self.keys is not None:
+            quant_cache.keys = mx.quantize(self.keys, group_size=group_size, bits=bits)
+            quant_cache.values = mx.quantize(
+                self.values, group_size=group_size, bits=bits
+            )
+        return quant_cache
+
+
+class RotatingKVCache(_BaseCache):
+
+    def __init__(self, max_size=None, keep=0, step=256):
+        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):
+        to_cat = []
+        if trim_size > 0:
+            to_cat = [v[..., : self.keep, :], v[..., trim_size + self.keep :, :]]
+        else:
+            to_cat = [v]
+        if append is not None:
+            to_cat.append(append)
+        return mx.concatenate(to_cat, axis=2)
+
+    def _temporal_order(self, v):
+        """
+        Rearrange the cache into temporal order, slicing off the end if unused.
+        """
+        if self._idx == v.shape[2]:
+            return v
+        elif self._idx < self.offset:
+            return mx.concatenate(
+                [
+                    v[..., : self.keep, :],
+                    v[..., self._idx :, :],
+                    v[..., self.keep : self._idx, :],
+                ],
+                axis=2,
+            )
+        else:
+            return v[..., : self._idx, :]
+
+    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.keys = self._temporal_order(self.keys)
+            self.values = self._temporal_order(self.values)
+
+            # The largest size is self.max_size + S to ensure
+            # every token gets at least self.max_size context
+            trim_size = self._idx - self.max_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._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
+
+        # Rotate
+        if self._idx == self.max_size:
+            self._idx = self.keep
+
+        # Assign
+        self.keys[..., self._idx : self._idx + S, :] = keys
+        self.values[..., self._idx : self._idx + S, :] = values
+        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):
+        if self.offset < self.keys.shape[2]:
+            return self.keys[..., : self.offset, :], self.values[..., : self.offset, :]
+        else:
+            return self.keys, self.values
+
+    @state.setter
+    def state(self, v):
+        self.keys, self.values = v
+
+    @property
+    def meta_state(self):
+        return tuple(
+            map(str, (self.keep, self.max_size, self.step, self.offset, self._idx))
+        )
+
+    @meta_state.setter
+    def meta_state(self, v):
+        self.keep, self.max_size, self.step, self.offset, self._idx = map(
+            int,
+            v,
+        )
+
+    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
+        return n
+
+    def to_quantized(self, group_size: int = 64, bits: int = 4) -> QuantizedKVCache:
+        raise NotImplementedError("RotatingKVCache Quantization NYI")
+
+
+class MambaCache(_BaseCache):
+    def __init__(self):
+        self.cache = [None, None]
+
+    def __setitem__(self, idx, value):
+        self.cache[idx] = value
+
+    def __getitem__(self, idx):
+        return self.cache[idx]
+
+    @property
+    def state(self):
+        return self.cache
+
+    @state.setter
+    def state(self, v):
+        self.cache = v
+
+
+class ChunkedKVCache(KVCache):
+    def __init__(self, chunk_size=None):
+        super().__init__()
+        self.chunk_size = chunk_size
+        self.start_position = 0
+
+    def maybe_trim_front(self):
+        # Maintain the cache below the chunk size
+        if self.keys is not None and self.keys.shape[2] >= self.chunk_size:
+            self.start_position += self.keys.shape[2] - self.chunk_size
+            self.keys = self.keys[..., -self.chunk_size :, :]
+            self.values = self.values[..., -self.chunk_size :, :]
+
+    def update_and_fetch(self, keys, values):
+        prev = self.offset - self.start_position
+        if self.keys is None or (prev + keys.shape[2]) > self.keys.shape[2]:
+            B, n_kv_heads, _, k_head_dim = keys.shape
+            v_head_dim = values.shape[3]
+            n_steps = (self.step + keys.shape[2] - 1) // self.step
+            k_shape = (B, n_kv_heads, n_steps * self.step, k_head_dim)
+            v_shape = (B, n_kv_heads, n_steps * self.step, v_head_dim)
+            new_k = mx.zeros(k_shape, keys.dtype)
+            new_v = mx.zeros(v_shape, values.dtype)
+            if self.keys is not None:
+                if prev % self.step != 0:
+                    self.keys = self.keys[..., :prev, :]
+                    self.values = self.values[..., :prev, :]
+                self.keys = mx.concatenate([self.keys, new_k], axis=2)
+                self.values = mx.concatenate([self.values, new_v], axis=2)
+            else:
+                self.keys, self.values = new_k, new_v
+
+        self.offset += keys.shape[2]
+        end = self.offset - self.start_position
+        self.keys[..., prev:end, :] = keys
+        self.values[..., prev:end, :] = values
+        return self.keys[..., :end, :], self.values[..., :end, :]
+
+    def trim(self, n):
+        n = min(self.offset - self.start_position, n)
+        self.offset -= n
+        return n
+
+    @property
+    def meta_state(self):
+        return tuple(map(str, (self.chunk_size, self.start_position)))
+
+    @meta_state.setter
+    def meta_state(self, v):
+        self.chunk_size, self.start_position = map(int, v)
+
+
+class CacheList(KVCache):
+    def __init__(self, *caches):
+        self.caches = caches
+
+    def __getitem__(self, idx):
+        return self.caches[idx]
+
+    @property
+    def state(self):
+        return [s for c in self.caches for s in c.state]
+
+    @state.setter
+    def state(self, v):
+        state_lens = [len(c.state) for c in self.caches]
+        start = 0
+        for c in self.caches:
+            l = len(c.state)
+            c.state = v[start : start + l]
+            start += l

mlx_lm/models/cohere.py

@@ -1,10 +1,12 @@
+# Copyright © 2023-2024 Apple Inc.
+
 from dataclasses import dataclass
-from typing import Optional, Tuple
+from typing import Any, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -67,7 +69,7 @@ class Attention(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         B, L, D = x.shape
 
@@ -91,8 +93,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
         )
 
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
@@ -127,7 +129,7 @@ class TransformerBlock(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         h = self.input_layernorm(x)
         attn_h = self.self_attn(h, mask, cache)
@@ -153,14 +155,13 @@ class CohereModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
         h = self.embed_tokens(inputs)
 
-        mask = None
-        if h.shape[1] > 1:
-            mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
-            mask = mask.astype(h.dtype)
+        if mask is None:
+            mask = create_attention_mask(h, cache)
 
         if cache is None:
             cache = [None] * len(self.layers)
@@ -181,9 +182,10 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, cache)
+        out = self.model(inputs, mask, cache)
         out = self.model.embed_tokens.as_linear(out)
         out = out * self.model.args.logit_scale
         return out
@@ -191,11 +193,3 @@ class Model(nn.Module):
     @property
     def layers(self):
         return self.model.layers
-
-    @property
-    def head_dim(self):
-        return self.args.hidden_size // self.args.num_attention_heads
-
-    @property
-    def n_kv_heads(self):
-        return self.args.num_key_value_heads

mlx_lm/models/cohere2.py

@@ -0,0 +1,227 @@
+# Copyright © 2023-2024 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import KVCache, RotatingKVCache
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int = 4096
+    head_dim: int = 128
+    num_hidden_layers: int = 32
+    intermediate_size: int = 14336
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 8
+    rope_theta: float = 50000.0
+    vocab_size: int = 256000
+    layer_norm_eps: float = 1e-05
+    logit_scale: float = 0.0625
+    attention_bias: bool = False
+    layer_norm_bias: bool = False
+    sliding_window: int = 4096
+    sliding_window_pattern: int = 4
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.args = args
+        self.layer_idx = layer_idx
+
+        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
+        if (head_dim * n_heads) != dim:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads (got `hidden_size`: {dim}"
+                f" and `num_heads`: {n_heads})."
+            )
+        self.scale = head_dim**-0.5
+
+        attetion_bias = args.attention_bias
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=attetion_bias)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attetion_bias)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attetion_bias)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=attetion_bias)
+
+        self.rope = nn.RoPE(head_dim, traditional=True, base=args.rope_theta)
+
+        self.use_sliding_window = (layer_idx + 1) % args.sliding_window_pattern != 0
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = 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 sliding window is enabled
+        if self.use_sliding_window:
+            if cache is None:
+                queries = self.rope(queries)
+                keys = self.rope(keys)
+            else:
+                queries = self.rope(queries, offset=cache.offset)
+                keys = self.rope(keys, offset=cache.offset)
+
+        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
+        output = scaled_dot_product_attention(
+            queries.astype(sdpa_type),
+            keys,
+            values,
+            cache=cache,
+            scale=self.scale,
+            mask=mask,
+        ).astype(queries.dtype)
+        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.up_proj = nn.Linear(dim, hidden_dim, bias=False)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
+
+    def __call__(self, x):
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.hidden_size = args.hidden_size
+        self.n_heads = args.num_attention_heads
+
+        self.self_attn = Attention(args, layer_idx)
+        self.mlp = MLP(args.hidden_size, args.intermediate_size)
+        self.input_layernorm = nn.LayerNorm(
+            args.hidden_size, eps=args.layer_norm_eps, bias=args.layer_norm_bias
+        )
+        self.args = args
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = None,
+    ) -> mx.array:
+
+        h = self.input_layernorm(x)
+        attn_h = self.self_attn(h, mask, cache)
+        ff_h = self.mlp(h)
+
+        return attn_h + ff_h + x
+
+
+class CohereModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.num_hidden_layers = args.num_hidden_layers
+        assert self.vocab_size > 0
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            TransformerBlock(args=args, layer_idx=i)
+            for i in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.LayerNorm(
+            args.hidden_size, eps=args.layer_norm_eps, bias=args.layer_norm_bias
+        )
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        h = self.embed_tokens(inputs)
+
+        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)
+
+        for i, (layer, c) in enumerate(zip(self.layers, cache)):
+            is_global = (
+                i % self.args.sliding_window_pattern
+                == 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
+
+            h = layer(h, local_mask, c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.model_type = args.model_type
+        self.model = CohereModel(args)
+        self.args = args
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        out = self.model(inputs, mask, cache)
+        out = self.model.embed_tokens.as_linear(out)
+        out = out * self.model.args.logit_scale
+        return out
+
+    def make_cache(self):
+        caches = []
+        for i in range(self.args.num_hidden_layers):
+            if (
+                i % self.args.sliding_window_pattern
+                == self.args.sliding_window_pattern - 1
+            ):
+                caches.append(KVCache())
+            else:
+                caches.append(
+                    RotatingKVCache(max_size=self.args.sliding_window, keep=0)
+                )
+        return caches
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/dbrx.py

@@ -1,11 +1,13 @@
+# Copyright © 2023-2024 Apple Inc.
+
 from dataclasses import dataclass
-from typing import Optional, Tuple
+from typing import Any, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
 import numpy as np
 
-from .base import BaseModelArgs
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -47,7 +49,7 @@ class Attention(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
 
         qkv = self.Wqkv(x)
@@ -72,8 +74,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
         )
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
         return self.out_proj(output)
@@ -90,7 +92,7 @@ class NormAttnNorm(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         h = self.attn(self.norm_1(x), mask=mask, cache=cache)
         x = h + x
@@ -103,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
 
@@ -177,7 +178,7 @@ class DecoderLayer(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         r, h = self.norm_attn_norm(x, mask, cache)
         out = self.ffn(h) + r
@@ -195,15 +196,13 @@ class DBRX(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
         h = self.wte(inputs)
 
-        mask = None
-        T = h.shape[1]
-        if T > 1:
-            mask = nn.MultiHeadAttention.create_additive_causal_mask(T)
-            mask = mask.astype(h.dtype)
+        if mask is None:
+            mask = create_attention_mask(h, cache)
 
         if cache is None:
             cache = [None] * len(self.blocks)
@@ -225,9 +224,10 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
-        out = self.transformer(inputs, cache)
+        out = self.transformer(inputs, mask, cache)
         return self.lm_head(out)
 
     @property
@@ -251,11 +251,3 @@ class Model(nn.Module):
                     experts = [(s, sv.T) for s, sv in experts]
                 new_weights.update(experts)
         return new_weights
-
-    @property
-    def head_dim(self):
-        return self.args.d_model // self.args.n_heads
-
-    @property
-    def n_kv_heads(self):
-        return self.args.attn_config["kv_n_heads"]

mlx_lm/models/deepseek.py

@@ -0,0 +1,260 @@
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str = "deepseek"
+    vocab_size: int = 102400
+    hidden_size: int = 4096
+    intermediate_size: int = 11008
+    moe_intermediate_size: int = 1407
+    num_hidden_layers: int = 30
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 32
+    n_shared_experts: Optional[int] = None
+    n_routed_experts: Optional[int] = None
+    num_experts_per_tok: Optional[int] = None
+    moe_layer_freq: int = 1
+    first_k_dense_replace: int = 0
+    max_position_embeddings: int = 2048
+    rms_norm_eps: float = 1e-6
+    rope_theta: float = 10000.0
+    rope_scaling: Optional[Dict] = None
+    attention_bias: bool = False
+
+
+class DeepseekAttention(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.num_attention_heads = config.num_attention_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.head_dim = config.hidden_size // config.num_attention_heads
+        self.scale = self.head_dim**-0.5
+
+        attention_bias = getattr(config, "attention_bias", False)
+
+        self.q_proj = nn.Linear(
+            self.hidden_size,
+            config.num_attention_heads * self.head_dim,
+            bias=attention_bias,
+        )
+        self.k_proj = nn.Linear(
+            self.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            self.hidden_size,
+            config.num_key_value_heads * self.head_dim,
+            bias=attention_bias,
+        )
+        self.o_proj = nn.Linear(
+            self.hidden_size,
+            config.num_attention_heads * self.head_dim,
+            bias=attention_bias,
+        )
+
+        rope_scale = 1.0
+        if config.rope_scaling and config.rope_scaling["type"] == "linear":
+            assert isinstance(config.rope_scaling["factor"], float)
+            rope_scale = 1 / config.rope_scaling["factor"]
+        self.rope = nn.RoPE(
+            self.head_dim,
+            base=config.rope_theta,
+            scale=rope_scale,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, _ = 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_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, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class DeepseekMLP(nn.Module):
+    def __init__(
+        self,
+        config: ModelArgs,
+        hidden_size: Optional[int] = None,
+        intermediate_size: Optional[int] = None,
+    ):
+        super().__init__()
+        self.config = config
+        self.hidden_size = hidden_size or config.hidden_size
+        self.intermediate_size = intermediate_size or 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)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class MoEGate(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.top_k = config.num_experts_per_tok
+        self.n_routed_experts = config.n_routed_experts
+        self.weight = mx.zeros((self.n_routed_experts, config.hidden_size))
+
+    def __call__(self, x):
+        gates = x @ self.weight.T
+        scores = mx.softmax(gates, axis=-1, precise=True)
+        k = self.top_k
+        inds = mx.stop_gradient(mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k])
+        scores = mx.take_along_axis(scores, inds, axis=-1)
+        return inds, scores
+
+
+class DeepseekMoE(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        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 = DeepseekMLP(
+                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)
+        if self.config.n_shared_experts is not None:
+            y = y + self.shared_experts(x)
+
+        return y
+
+
+class DeepseekDecoderLayer(nn.Module):
+    def __init__(self, config: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.self_attn = DeepseekAttention(config)
+        self.mlp = (
+            DeepseekMoE(config)
+            if (
+                config.n_routed_experts is not None
+                and layer_idx >= config.first_k_dense_replace
+                and layer_idx % config.moe_layer_freq == 0
+            )
+            else DeepseekMLP(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))
+        out = h + r
+        return out
+
+
+class DeepseekModel(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = [
+            DeepseekDecoderLayer(config, idx) for idx in range(config.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def __call__(
+        self,
+        x: mx.array,
+        cache: Optional[Any] = None,
+        mask: Optional[mx.array] = None,
+    ) -> mx.array:
+        h = self.embed_tokens(x)
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.args = config
+        self.model_type = config.model_type
+        self.model = DeepseekModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+        mask: Optional[mx.array] = None,
+    ):
+        out = self.model(inputs, cache, mask)
+        return self.lm_head(out)
+
+    def sanitize(self, weights):
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}"
+            for m in ["gate_proj", "down_proj", "up_proj"]:
+                for k in ["weight", "scales", "biases"]:
+                    if f"{prefix}.mlp.experts.0.{m}.{k}" in weights:
+                        to_join = [
+                            weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}")
+                            for e in range(self.args.n_routed_experts)
+                        ]
+                        weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
+        return weights
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/deepseek_v2.py

@@ -0,0 +1,458 @@
+# Copyright © 2023-2024 Apple Inc.
+
+import math
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str = "deepseek_v2"
+    vocab_size: int = 102400
+    hidden_size: int = 4096
+    intermediate_size: int = 11008
+    moe_intermediate_size: int = 1407
+    num_hidden_layers: int = 30
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 32
+    n_shared_experts: Optional[int] = None
+    n_routed_experts: Optional[int] = None
+    routed_scaling_factor: float = 1.0
+    kv_lora_rank: int = 512
+    q_lora_rank: int = 1536
+    qk_rope_head_dim: int = 64
+    v_head_dim: int = 128
+    qk_nope_head_dim: int = 128
+    topk_method: str = "gready"
+    n_group: Optional[int] = None
+    topk_group: Optional[int] = None
+    num_experts_per_tok: Optional[int] = None
+    moe_layer_freq: int = 1
+    first_k_dense_replace: int = 0
+    max_position_embeddings: int = 2048
+    rms_norm_eps: float = 1e-6
+    rope_theta: float = 10000.0
+    rope_scaling: Dict = None
+    attention_bias: bool = False
+
+
+def yarn_find_correction_dim(
+    num_rotations, dim, base=10000, max_position_embeddings=2048
+):
+    return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (
+        2 * math.log(base)
+    )
+
+
+def yarn_find_correction_range(
+    low_rot, high_rot, dim, base=10000, max_position_embeddings=2048
+):
+    low = math.floor(
+        yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings)
+    )
+    high = math.ceil(
+        yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings)
+    )
+    return max(low, 0), min(high, dim - 1)
+
+
+def yarn_get_mscale(scale=1, mscale=1):
+    if scale <= 1:
+        return 1.0
+    return 0.1 * mscale * math.log(scale) + 1.0
+
+
+def yarn_linear_ramp_mask(min_val, max_val, dim):
+    if min_val == max_val:
+        max_val += 0.001  # Prevent singularity
+
+    linear_func = (mx.arange(dim, dtype=mx.float32) - min_val) / (max_val - min_val)
+    return mx.clip(linear_func, 0, 1)
+
+
+class DeepseekV2YarnRotaryEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim,
+        max_position_embeddings=2048,
+        base=10000,
+        scaling_factor=1.0,
+        original_max_position_embeddings=4096,
+        beta_fast=32,
+        beta_slow=1,
+        mscale=1,
+        mscale_all_dim=0,
+    ):
+        super().__init__()
+        self.mscale = yarn_get_mscale(scaling_factor, mscale) / yarn_get_mscale(
+            scaling_factor, mscale_all_dim
+        )
+        freq_extra = base ** (mx.arange(0, dim, 2, dtype=mx.float32) / dim)
+        freq_inter = scaling_factor * base ** (
+            mx.arange(0, dim, 2, dtype=mx.float32) / dim
+        )
+        low, high = yarn_find_correction_range(
+            beta_fast,
+            beta_slow,
+            dim,
+            base,
+            original_max_position_embeddings,
+        )
+        freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dim // 2)
+        self._freqs = (freq_inter * freq_extra) / (
+            freq_inter * freq_mask + freq_extra * (1 - freq_mask)
+        )
+
+    def __call__(self, x, offset=0):
+        if self.mscale != 1.0:
+            x = self.mscale * x
+        return mx.fast.rope(
+            x,
+            x.shape[-1],
+            traditional=True,
+            base=None,
+            scale=1.0,
+            offset=offset,
+            freqs=self._freqs,
+        )
+
+
+class DeepseekV2Attention(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rope_theta = config.rope_theta
+        self.q_lora_rank = config.q_lora_rank
+        self.qk_rope_head_dim = config.qk_rope_head_dim
+        self.kv_lora_rank = config.kv_lora_rank
+        self.v_head_dim = config.v_head_dim
+        self.qk_nope_head_dim = config.qk_nope_head_dim
+        self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim
+
+        self.scale = self.q_head_dim**-0.5
+
+        if self.q_lora_rank is None:
+            self.q_proj = nn.Linear(
+                self.hidden_size, self.num_heads * self.q_head_dim, bias=False
+            )
+        else:
+            self.q_a_proj = nn.Linear(
+                self.hidden_size, self.q_lora_rank, bias=config.attention_bias
+            )
+            self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank, eps=1e-6)
+            self.q_b_proj = nn.Linear(
+                self.q_lora_rank, self.num_heads * self.q_head_dim, bias=False
+            )
+
+        self.kv_a_proj_with_mqa = nn.Linear(
+            self.hidden_size,
+            self.kv_lora_rank + self.qk_rope_head_dim,
+            bias=config.attention_bias,
+        )
+        self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank, eps=1e-6)
+        self.kv_b_proj = nn.Linear(
+            self.kv_lora_rank,
+            self.num_heads
+            * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
+            bias=False,
+        )
+
+        self.o_proj = nn.Linear(
+            self.num_heads * self.v_head_dim,
+            self.hidden_size,
+            bias=config.attention_bias,
+        )
+
+        mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0)
+        scaling_factor = self.config.rope_scaling["factor"]
+        if mscale_all_dim:
+            mscale = yarn_get_mscale(scaling_factor, mscale_all_dim)
+            self.scale = self.scale * mscale * mscale
+
+        rope_kwargs = {
+            key: self.config.rope_scaling[key]
+            for key in [
+                "original_max_position_embeddings",
+                "beta_fast",
+                "beta_slow",
+                "mscale",
+                "mscale_all_dim",
+            ]
+            if key in self.config.rope_scaling
+        }
+        self.rope = DeepseekV2YarnRotaryEmbedding(
+            dim=self.qk_rope_head_dim,
+            max_position_embeddings=self.max_position_embeddings,
+            scaling_factor=scaling_factor,
+            base=self.rope_theta,
+            **rope_kwargs,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        if self.q_lora_rank is None:
+            q = self.q_proj(x)
+        else:
+            q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x)))
+
+        q = q.reshape(B, L, self.num_heads, self.q_head_dim).transpose(0, 2, 1, 3)
+        q_nope, q_pe = mx.split(q, [self.qk_nope_head_dim], axis=-1)
+        compressed_kv = self.kv_a_proj_with_mqa(x)
+        compressed_kv, k_pe = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
+        k_pe = k_pe.reshape(B, L, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3)
+        kv = self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
+        kv = kv.reshape(B, L, self.num_heads, -1).transpose(0, 2, 1, 3)
+
+        k_nope, values = mx.split(kv, [self.qk_nope_head_dim], axis=-1)
+
+        if cache is not None:
+            q_pe = self.rope(q_pe, cache.offset)
+            k_pe = self.rope(k_pe, cache.offset)
+            k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
+            keys, values = cache.update_and_fetch(
+                mx.concatenate([k_nope, k_pe], axis=-1), values
+            )
+        else:
+            q_pe = self.rope(q_pe)
+            k_pe = self.rope(k_pe)
+            k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
+            keys = mx.concatenate([k_nope, k_pe], axis=-1)
+
+        queries = mx.concatenate([q_nope, q_pe], axis=-1)
+
+        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 DeepseekV2MLP(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
+
+
+class MoEGate(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.top_k = config.num_experts_per_tok
+        self.n_routed_experts = config.n_routed_experts
+        self.routed_scaling_factor = config.routed_scaling_factor
+        self.topk_method = config.topk_method
+        self.n_group = config.n_group
+        self.topk_group = config.topk_group
+        self.weight = mx.zeros((self.n_routed_experts, config.hidden_size))
+
+    def __call__(self, x):
+        gates = x @ self.weight.T
+
+        scores = mx.softmax(gates, axis=-1, precise=True)
+
+        if self.topk_method == "group_limited_greedy":
+            bsz, seq_len = x.shape[:2]
+            scores = scores.reshape(bsz, seq_len, self.n_group, -1)
+            group_scores = scores.max(axis=-1, keepdims=True)
+            k = self.n_group - self.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, scores.dtype), axis=-2
+            )
+            scores = scores.reshape(bsz, seq_len, -1)
+
+        k = self.top_k
+        inds = mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k]
+        scores = mx.take_along_axis(scores, inds, axis=-1)
+        scores = scores * self.routed_scaling_factor
+
+        return inds, scores
+
+
+class DeepseekV2MoE(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 = DeepseekV2MLP(
+                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)
+        if self.config.n_shared_experts is not None:
+            y = y + self.shared_experts(x)
+
+        return y
+
+
+class DeepseekV2DecoderLayer(nn.Module):
+    def __init__(self, config: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.self_attn = DeepseekV2Attention(config)
+        self.mlp = (
+            DeepseekV2MoE(config)
+            if (
+                config.n_routed_experts is not None
+                and layer_idx >= config.first_k_dense_replace
+                and layer_idx % config.moe_layer_freq == 0
+            )
+            else DeepseekV2MLP(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))
+        out = h + r
+        return out
+
+
+class DeepseekV2Model(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 = [
+            DeepseekV2DecoderLayer(config, idx)
+            for idx in range(config.num_hidden_layers)
+        ]
+        self.start_idx = 0
+        self.end_idx = len(self.layers)
+        self.num_layers = self.end_idx
+
+        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        self.pipeline_rank = 0
+        self.pipeline_size = 1
+
+    def pipeline(self, group):
+        # Split layers in reverse so rank=0 gets the last layers and
+        # rank=pipeline_size-1 gets the first
+        self.pipeline_rank = group.rank()
+        self.pipeline_size = group.size()
+        layers_per_rank = len(self.layers) // self.pipeline_size
+        extra = len(self.layers) - layers_per_rank * self.pipeline_size
+        if self.pipeline_rank < extra:
+            layers_per_rank += 1
+
+        self.start_idx = (self.pipeline_size - self.pipeline_rank - 1) * layers_per_rank
+        self.end_idx = self.start_idx + layers_per_rank
+        self.num_layers = layers_per_rank
+        self.layers = self.layers[: self.end_idx]
+        self.layers[: self.start_idx] = [None] * self.start_idx
+        self.num_layers = len(self.layers) - self.start_idx
+
+    def __call__(
+        self,
+        x: mx.array,
+        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
+
+        # Receive from the previous process in the pipeline
+        if pipeline_rank < pipeline_size - 1:
+            h = mx.distributed.recv_like(h, (pipeline_rank + 1))
+
+        for i in range(self.num_layers):
+            h = self.layers[self.start_idx + i](h, mask, cache[i])
+
+        # Send to the next process in the pipeline
+        if pipeline_rank != 0:
+            h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
+
+        # Broadcast h while keeping it in the graph
+        h = mx.distributed.all_gather(h)[: h.shape[0]]
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.args = config
+        self.model_type = config.model_type
+        self.model = DeepseekV2Model(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+        mask: Optional[mx.array] = None,
+    ):
+        out = self.model(inputs, cache, mask)
+        return self.lm_head(out)
+
+    def sanitize(self, weights):
+        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)
+        return weights
+
+    @property
+    def layers(self):
+        return self.model.layers[self.model.start_idx : self.model.end_idx]

mlx_lm/models/deepseek_v3.py

@@ -0,0 +1,545 @@
+# Copyright © 2024 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 = "deepseek_v3"
+    vocab_size: int = 102400
+    hidden_size: int = 4096
+    intermediate_size: int = 11008
+    moe_intermediate_size: int = 1407
+    num_hidden_layers: int = 30
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 32
+    n_shared_experts: Optional[int] = None
+    n_routed_experts: Optional[int] = None
+    routed_scaling_factor: float = 1.0
+    kv_lora_rank: int = 512
+    q_lora_rank: int = 1536
+    qk_rope_head_dim: int = 64
+    v_head_dim: int = 128
+    qk_nope_head_dim: int = 128
+    topk_method: str = "noaux_tc"
+    scoring_func: str = "sigmoid"
+    norm_topk_prob: bool = True
+    n_group: Optional[int] = None
+    topk_group: Optional[int] = None
+    num_experts_per_tok: Optional[int] = None
+    moe_layer_freq: int = 1
+    first_k_dense_replace: int = 0
+    max_position_embeddings: int = 2048
+    rms_norm_eps: float = 1e-6
+    rope_theta: float = 10000.0
+    rope_scaling: Dict = None
+    attention_bias: bool = False
+
+
+def yarn_find_correction_dim(
+    num_rotations, dim, base=10000, max_position_embeddings=2048
+):
+    return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (
+        2 * math.log(base)
+    )
+
+
+def yarn_find_correction_range(
+    low_rot, high_rot, dim, base=10000, max_position_embeddings=2048
+):
+    low = math.floor(
+        yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings)
+    )
+    high = math.ceil(
+        yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings)
+    )
+    return max(low, 0), min(high, dim - 1)
+
+
+def yarn_get_mscale(scale=1, mscale=1):
+    if scale <= 1:
+        return 1.0
+    return 0.1 * mscale * math.log(scale) + 1.0
+
+
+def yarn_linear_ramp_mask(min_val, max_val, dim):
+    if min_val == max_val:
+        max_val += 0.001  # Prevent singularity
+
+    linear_func = (mx.arange(dim, dtype=mx.float32) - min_val) / (max_val - min_val)
+    return mx.clip(linear_func, 0, 1)
+
+
+class DeepseekV3YarnRotaryEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim,
+        max_position_embeddings=2048,
+        base=10000,
+        scaling_factor=1.0,
+        original_max_position_embeddings=4096,
+        beta_fast=32,
+        beta_slow=1,
+        mscale=1,
+        mscale_all_dim=0,
+    ):
+        super().__init__()
+        self.mscale = yarn_get_mscale(scaling_factor, mscale) / yarn_get_mscale(
+            scaling_factor, mscale_all_dim
+        )
+        freq_extra = base ** (mx.arange(0, dim, 2, dtype=mx.float32) / dim)
+        freq_inter = scaling_factor * freq_extra
+        low, high = yarn_find_correction_range(
+            beta_fast,
+            beta_slow,
+            dim,
+            base,
+            original_max_position_embeddings,
+        )
+        freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dim // 2)
+        self._freqs = (freq_inter * freq_extra) / (
+            freq_inter * freq_mask + freq_extra * (1 - freq_mask)
+        )
+
+    def __call__(self, x, offset=0):
+        if self.mscale != 1.0:
+            x = self.mscale * x
+        return mx.fast.rope(
+            x,
+            x.shape[-1],
+            traditional=True,
+            base=None,
+            scale=1.0,
+            offset=offset,
+            freqs=self._freqs,
+        )
+
+
+# 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 ClippedSilu(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def __call__(self, x):
+        return clipped_silu(x)
+
+
+class DeepseekV3Attention(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.max_position_embeddings = config.max_position_embeddings
+        self.rope_theta = config.rope_theta
+        self.q_lora_rank = config.q_lora_rank
+        self.qk_rope_head_dim = config.qk_rope_head_dim
+        self.kv_lora_rank = config.kv_lora_rank
+        self.v_head_dim = config.v_head_dim
+        self.qk_nope_head_dim = config.qk_nope_head_dim
+        self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim
+
+        self.scale = self.q_head_dim**-0.5
+
+        if self.q_lora_rank is None:
+            self.q_proj = nn.Linear(
+                self.hidden_size, self.num_heads * self.q_head_dim, bias=False
+            )
+        else:
+            self.q_a_proj = nn.Linear(
+                self.hidden_size, self.q_lora_rank, bias=config.attention_bias
+            )
+            self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank, eps=1e-6)
+            self.q_b_proj = nn.Linear(
+                self.q_lora_rank, self.num_heads * self.q_head_dim, bias=False
+            )
+
+        self.kv_a_proj_with_mqa = nn.Linear(
+            self.hidden_size,
+            self.kv_lora_rank + self.qk_rope_head_dim,
+            bias=config.attention_bias,
+        )
+        self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank, eps=1e-6)
+        self.kv_b_proj = nn.Linear(
+            self.kv_lora_rank,
+            self.num_heads
+            * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
+            bias=False,
+        )
+
+        self.o_proj = nn.Linear(
+            self.num_heads * self.v_head_dim,
+            self.hidden_size,
+            bias=config.attention_bias,
+        )
+
+        if self.config.rope_scaling is not None:
+            mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0)
+            scaling_factor = self.config.rope_scaling["factor"]
+            if mscale_all_dim:
+                mscale = yarn_get_mscale(scaling_factor, mscale_all_dim)
+                self.scale = self.scale * mscale * mscale
+
+            rope_kwargs = {
+                key: self.config.rope_scaling[key]
+                for key in [
+                    "original_max_position_embeddings",
+                    "beta_fast",
+                    "beta_slow",
+                    "mscale",
+                    "mscale_all_dim",
+                ]
+                if key in self.config.rope_scaling
+            }
+            self.rope = DeepseekV3YarnRotaryEmbedding(
+                dim=self.qk_rope_head_dim,
+                max_position_embeddings=self.max_position_embeddings,
+                scaling_factor=scaling_factor,
+                base=self.rope_theta,
+                **rope_kwargs,
+            )
+        else:
+            self.rope = nn.RoPE(
+                dims=self.qk_rope_head_dim,
+                base=self.rope_theta,
+                traditional=True,
+            )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        if self.q_lora_rank is None:
+            q = self.q_proj(x)
+        else:
+            q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x)))
+
+        q = q.reshape(B, L, self.num_heads, self.q_head_dim).transpose(0, 2, 1, 3)
+        q_nope, q_pe = mx.split(q, [self.qk_nope_head_dim], axis=-1)
+        compressed_kv = self.kv_a_proj_with_mqa(x)
+        compressed_kv, k_pe = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
+        k_pe = k_pe.reshape(B, L, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3)
+        kv = self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
+        kv = kv.reshape(B, L, self.num_heads, -1).transpose(0, 2, 1, 3)
+
+        k_nope, values = mx.split(kv, [self.qk_nope_head_dim], axis=-1)
+
+        if cache is not None:
+            q_pe = self.rope(q_pe, cache.offset)
+            k_pe = self.rope(k_pe, cache.offset)
+            k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
+            keys, values = cache.update_and_fetch(
+                mx.concatenate([k_nope, k_pe], axis=-1), values
+            )
+        else:
+            q_pe = self.rope(q_pe)
+            k_pe = self.rope(k_pe)
+            k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
+            keys = mx.concatenate([k_nope, k_pe], axis=-1)
+
+        queries = mx.concatenate([q_nope, q_pe], axis=-1)
+
+        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 DeepseekV3MLP(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,
+):
+
+    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)
+
+    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 DeepseekV3MoE(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,
+            activation=ClippedSilu(),
+        )
+
+        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 = DeepseekV3MLP(
+                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 DeepseekV3DecoderLayer(nn.Module):
+    def __init__(self, config: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.self_attn = DeepseekV3Attention(config)
+        self.mlp = (
+            DeepseekV3MoE(config)
+            if (
+                config.n_routed_experts is not None
+                and layer_idx >= config.first_k_dense_replace
+                and layer_idx % config.moe_layer_freq == 0
+            )
+            else DeepseekV3MLP(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 DeepseekV3Model(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 = [
+            DeepseekV3DecoderLayer(config, idx)
+            for idx in range(config.num_hidden_layers)
+        ]
+        self.start_idx = 0
+        self.end_idx = len(self.layers)
+        self.num_layers = self.end_idx
+
+        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.pipeline_rank = 0
+        self.pipeline_size = 1
+
+    def pipeline(self, group):
+        # Split layers in reverse so rank=0 gets the last layers and
+        # rank=pipeline_size-1 gets the first
+        self.pipeline_rank = group.rank()
+        self.pipeline_size = group.size()
+        layers_per_rank = len(self.layers) // self.pipeline_size
+        extra = len(self.layers) - layers_per_rank * self.pipeline_size
+        if self.pipeline_rank < extra:
+            layers_per_rank += 1
+        self.start_idx = (self.pipeline_size - self.pipeline_rank - 1) * layers_per_rank
+        self.end_idx = self.start_idx + layers_per_rank
+        self.layers = self.layers[: self.end_idx]
+        self.layers[: self.start_idx] = [None] * self.start_idx
+        self.num_layers = len(self.layers) - self.start_idx
+
+    def __call__(
+        self,
+        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
+
+        # Receive from the previous process in the pipeline
+
+        if pipeline_rank < pipeline_size - 1:
+            h = mx.distributed.recv_like(h, (pipeline_rank + 1))
+
+        for i in range(self.num_layers):
+            h = self.layers[self.start_idx + i](h, mask, cache[i])
+
+        # Send to the next process in the pipeline
+        if pipeline_rank != 0:
+            h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
+
+        # Broadcast h while keeping it in the graph
+        h = mx.distributed.all_gather(h)[: h.shape[0]]
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.args = config
+        self.model_type = config.model_type
+        self.model = DeepseekV3Model(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+        mask: Optional[mx.array] = None,
+    ):
+        out = self.model(inputs, cache, mask)
+        return self.lm_head(out)
+
+    def sanitize(self, weights):
+        def dequant(weight, scale_inv):
+            dtype = weight.dtype
+            bs = 128  # block size
+            m, n = weight.shape
+            pad_bottom = (-m) % bs
+            pad_side = (-n) % bs
+            weight = mx.pad(weight, ((0, pad_bottom), (0, pad_side)))
+            weight = weight.reshape(
+                ((m + pad_bottom) // bs, bs, (n + pad_side) // bs, bs)
+            )
+            weight = (weight * scale_inv[:, None, :, None]).reshape(
+                m + pad_bottom, n + pad_side
+            )
+            return weight[:m, :n].astype(dtype)
+
+        # Dequantize
+        new_weights = {}
+        for k, v in weights.items():
+            if "weight_scale_inv" in k:
+                scale_inv = v
+                wk = k.replace("_scale_inv", "")
+                weight = weights[wk]
+                weight = dequant(weight, scale_inv)
+                new_weights[wk] = weight
+            elif k not in new_weights:
+                new_weights[k] = v
+        weights = new_weights
+
+        # Stack experts
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}"
+            for n, m in [("w1", "gate_proj"), ("w2", "down_proj"), ("w3", "up_proj")]:
+                for k in ["weight", "scales", "biases"]:
+                    if f"{prefix}.mlp.experts.0.{m}.{k}" in weights:
+                        to_join = [
+                            weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}")
+                            for e in range(self.args.n_routed_experts)
+                        ]
+                        weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
+
+        # Remove multi-token prediction layer and any unused precomputed rotary freqs
+        return {
+            k: v
+            for k, v in weights.items()
+            if not k.startswith("model.layers.61") and "rotary_emb.inv_freq" not in k
+        }
+
+    @property
+    def layers(self):
+        return self.model.layers[self.model.start_idx : self.model.end_idx]
+
+    @property
+    def cast_predicate(self):
+        def predicate(k):
+            return "e_score_correction_bias" not in k
+
+        return predicate

mlx_lm/models/dots1.py

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

mlx_lm/models/ernie4_5.py

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

mlx_lm/models/ernie4_5_moe.py

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

mlx_lm/models/exaone.py

@@ -0,0 +1,166 @@
+# Copyright © 2024 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    vocab_size: int
+    rope_theta: float
+    layer_norm_epsilon: float
+    num_key_value_heads: int
+    head_dim: Optional[int] = None
+    max_position_embeddings: Optional[int] = None
+    rope_traditional: bool = False
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = True
+    attention_bias: bool = False
+    mlp_bias: bool = False
+
+
+class AttentionModule(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.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.out_proj = nn.Linear(n_heads * head_dim, dim, 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: mx.array, mask: Optional[mx.array] = None, cache: Optional[Any] = None
+    ) -> mx.array:
+        B, L, D = x.shape
+        q = self.q_proj(x).reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        k = self.k_proj(x).reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        v = self.v_proj(x).reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        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)
+
+        out = scaled_dot_product_attention(
+            q, k, v, cache=cache, scale=self.scale, mask=mask
+        )
+        out = out.transpose(0, 2, 1, 3).reshape(B, L, D)
+        return self.out_proj(out)
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.attention = AttentionModule(args)
+
+
+class MLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        dim = args.hidden_size
+        hidden_dim = args.intermediate_size
+        self.c_fc_0 = nn.Linear(dim, hidden_dim, bias=args.mlp_bias)
+        self.c_fc_1 = nn.Linear(dim, hidden_dim, bias=args.mlp_bias)
+        self.c_proj = nn.Linear(hidden_dim, dim, bias=args.mlp_bias)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.c_proj(nn.silu(self.c_fc_0(x)) * self.c_fc_1(x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.ln_1 = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon)
+        self.attn = Attention(args)
+        self.ln_2 = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon)
+        self.mlp = MLP(args)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        h = x + self.attn.attention(self.ln_1(x), mask, cache)
+        out = h + self.mlp(self.ln_2(h))
+        return out
+
+
+class ExaoneModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.wte = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.h = [TransformerBlock(args) for _ in range(args.num_layers)]
+        self.ln_f = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon)
+
+    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)
+
+        for layer, c in zip(self.h, cache):
+            h = layer(h, mask, cache=c)
+
+        return self.ln_f(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.transformer = ExaoneModel(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        out = self.transformer(inputs, mask, cache)
+        if self.args.tie_word_embeddings:
+            out = self.transformer.wte.as_linear(out)
+        else:
+            out = self.lm_head(out)
+        return out
+
+    @property
+    def layers(self):
+        return self.transformer.h

mlx_lm/models/gemma.py

@@ -1,10 +1,12 @@
+# Copyright © 2023-2024 Apple Inc.
+
 from dataclasses import dataclass
-from typing import Optional, Tuple
+from typing import Any, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -58,7 +60,7 @@ class Attention(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         B, L, D = x.shape
 
@@ -77,8 +79,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
         )
 
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
@@ -111,7 +113,7 @@ class TransformerBlock(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         r = self.self_attn(self.input_layernorm(x), mask, cache)
         h = x + r
@@ -136,15 +138,14 @@ 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)
 
-        mask = None
-        if h.shape[1] > 1:
-            mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
-            mask = mask.astype(h.dtype)
+        if mask is None:
+            mask = create_attention_mask(h, cache)
 
         if cache is None:
             cache = [None] * len(self.layers)
@@ -165,20 +166,13 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, cache)
+        out = self.model(inputs, mask, cache)
         out = self.model.embed_tokens.as_linear(out)
         return out
 
     @property
     def layers(self):
         return self.model.layers
-
-    @property
-    def head_dim(self):
-        return self.args.head_dim
-
-    @property
-    def n_kv_heads(self):
-        return self.args.num_key_value_heads

mlx_lm/models/gemma2.py

@@ -0,0 +1,208 @@
+# 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
+
+
+@dataclass
+class ModelArgs(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
+    rope_theta: float = 10000
+    rope_traditional: bool = False
+    attn_logit_softcapping: float = 50.0
+    final_logit_softcapping: float = 30.0
+    query_pre_attn_scalar: float = 144.0
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dims: int, eps: float = 1e-5):
+        super().__init__()
+        self.weight = mx.ones((dims,))
+        self.eps = eps
+
+    def __call__(self, x):
+        return mx.fast.rms_norm(x, 1.0 + self.weight, self.eps)
+
+
+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.repeats = n_heads // n_kv_heads
+        self.head_dim = head_dim = args.head_dim
+
+        self.scale = 1.0 / (args.query_pre_attn_scalar**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.attn_logit_softcapping = args.attn_logit_softcapping
+        self.rope = nn.RoPE(
+            head_dim,
+            traditional=args.rope_traditional,
+            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).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)
+
+        queries = queries * self.scale
+
+        if self.repeats > 1:
+            queries = queries.reshape(
+                B, self.n_kv_heads, self.repeats, L, self.head_dim
+            )
+            keys = mx.expand_dims(keys, 2)
+            values = mx.expand_dims(values, 2)
+
+        scores = queries @ keys.swapaxes(-1, -2)
+        scores = mx.tanh(scores / self.attn_logit_softcapping)
+        scores *= self.attn_logit_softcapping
+
+        if mask is not None:
+            if mask.dtype == mx.bool_:
+                scores = mx.where(
+                    mask, scores, mx.array(mx.finfo(scores.dtype).min, scores.dtype)
+                )
+            else:
+                scores = scores + mask
+        scores = mx.softmax(scores, precise=True, axis=-1)
+        output = scores @ values
+        if self.repeats > 1:
+            output = output.reshape(B, self.n_heads, L, self.head_dim)
+        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.gelu_approx(self.gate_proj(x)) * self.up_proj(x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.hidden_size = args.hidden_size
+        self.self_attn = Attention(args)
+        self.mlp = MLP(args.hidden_size, args.intermediate_size)
+        self.input_layernorm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.pre_feedforward_layernorm = RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+        self.post_feedforward_layernorm = RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+        self.post_attention_layernorm = 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 + self.post_attention_layernorm(r)
+        r = self.mlp(self.pre_feedforward_layernorm(h))
+        out = h + self.post_feedforward_layernorm(r)
+        return out
+
+
+class GemmaModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.num_hidden_layers = args.num_hidden_layers
+        assert self.vocab_size > 0
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            TransformerBlock(args=args) for _ in range(args.num_hidden_layers)
+        ]
+        self.norm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        h = self.embed_tokens(inputs)
+        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)
+
+        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.model_type = args.model_type
+        self.final_logit_softcapping = args.final_logit_softcapping
+        self.model = GemmaModel(args)
+        self.args = args
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        out = self.model(inputs, mask, cache)
+        out = self.model.embed_tokens.as_linear(out)
+        out = mx.tanh(out / self.final_logit_softcapping)
+        out = out * self.final_logit_softcapping
+        return out
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/gemma3.py

@@ -0,0 +1,64 @@
+# 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 gemma3_text
+from .base import BaseModelArgs
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    text_config: dict
+    vocab_size: int = 262208
+
+    def __post_init__(self):
+        self.text_config["vocab_size"] = self.vocab_size
+        self.text_config["num_attention_heads"] = self.text_config.get(
+            "num_attention_heads", 8
+        )
+        self.text_config["num_key_value_heads"] = self.text_config.get(
+            "num_key_value_heads", 4
+        )
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.language_model = gemma3_text.Model(
+            gemma3_text.ModelArgs.from_dict(args.text_config)
+        )
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+        mask: Optional[mx.array] = None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        return self.language_model(
+            inputs, cache=cache, mask=mask, input_embeddings=input_embeddings
+        )
+
+    def sanitize(self, weights):
+        weights = tree_unflatten(list(weights.items()))
+        weights.pop("vision_tower", None)
+        weights.pop("multi_modal_projector", None)
+        lm_weights = dict(tree_flatten(weights["language_model"]))
+        lm_weights = self.language_model.sanitize(lm_weights)
+        weights["language_model"] = tree_unflatten(list(lm_weights.items()))
+        return dict(tree_flatten(weights))
+
+    @property
+    def layers(self):
+        return self.language_model.layers
+
+    def make_cache(self):
+        return self.language_model.make_cache()

mlx_lm/models/gemma3_text.py

@@ -0,0 +1,252 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import KVCache, RotatingKVCache
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int = 1152
+    num_hidden_layers: int = 26
+    intermediate_size: int = 6912
+    num_attention_heads: int = 4
+    head_dim: int = 256
+    rms_norm_eps: float = 1.0e-6
+    vocab_size: int = 262144
+    num_key_value_heads: int = 1
+    rope_global_base_freq: float = 1_000_000.0
+    rope_local_base_freq: float = 10_000.0
+    rope_traditional: bool = False
+    query_pre_attn_scalar: float = 256
+    sliding_window: int = 512
+    sliding_window_pattern: int = 6
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+        self.repeats = n_heads // n_kv_heads
+        self.head_dim = head_dim = args.head_dim
+        self.layer_idx = layer_idx
+
+        self.scale = args.query_pre_attn_scalar**-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 = RMSNorm(dims=head_dim, eps=args.rms_norm_eps)
+        self.k_norm = RMSNorm(dims=head_dim, eps=args.rms_norm_eps)
+        self.is_sliding = (layer_idx + 1) % args.sliding_window_pattern != 0
+
+        self.rope = nn.RoPE(
+            head_dim,
+            traditional=args.rope_traditional,
+            base=(
+                args.rope_local_base_freq
+                if self.is_sliding
+                else args.rope_global_base_freq
+            ),
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, _ = 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)
+
+        queries = self.q_norm(queries)
+        keys = self.k_norm(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)
+
+        # 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
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dims: int, eps: float = 1e-5):
+        super().__init__()
+        self.weight = mx.ones((dims,))
+        self.eps = eps
+
+    def __call__(self, x):
+        return mx.fast.rms_norm(x, 1.0 + self.weight, self.eps)
+
+
+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.gelu_approx(self.gate_proj(x)) * self.up_proj(x))
+
+
+@partial(mx.compile, shapeless=True)
+def clip_residual(x, y):
+    if x.dtype != mx.float16:
+        return x + y
+    bound = mx.finfo(mx.float16).max
+    return mx.clip(x.astype(mx.float32) + y.astype(mx.float32), -bound, bound).astype(
+        mx.float16
+    )
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.hidden_size = args.hidden_size
+        self.self_attn = Attention(args, layer_idx)
+        self.mlp = MLP(args.hidden_size, args.intermediate_size)
+        self.input_layernorm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.pre_feedforward_layernorm = RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+        self.post_feedforward_layernorm = 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 = clip_residual(x, self.post_attention_layernorm(r))
+        r = self.mlp(self.pre_feedforward_layernorm(h))
+        out = clip_residual(h, self.post_feedforward_layernorm(r))
+        return out
+
+
+class Gemma3Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.num_hidden_layers = args.num_hidden_layers
+        assert self.vocab_size > 0
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            TransformerBlock(args=args, layer_idx=layer_idx)
+            for layer_idx in range(args.num_hidden_layers)
+        ]
+        self.norm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+        input_embeddings: Optional[mx.array] = None,
+    ):
+        if input_embeddings is not None:
+            h = input_embeddings
+        else:
+            h = self.embed_tokens(inputs)
+        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)
+
+        for i, (layer, c) in enumerate(zip(self.layers, cache)):
+            is_global = (
+                i % self.args.sliding_window_pattern
+                == 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
+
+            h = layer(h, local_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 = Gemma3Model(args)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=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, input_embeddings)
+        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"]
+        return weights
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def make_cache(self):
+        caches = []
+        for i in range(self.args.num_hidden_layers):
+            if (
+                i % self.args.sliding_window_pattern
+                == self.args.sliding_window_pattern - 1
+            ):
+                caches.append(KVCache())
+            else:
+                caches.append(
+                    RotatingKVCache(max_size=self.args.sliding_window, keep=0)
+                )
+        return caches

mlx_lm/models/gemma3n.py

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

mlx_lm/models/glm4.py

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

mlx_lm/models/gpt2.py

@@ -1,11 +1,12 @@
+# Copyright © 2023-2024 Apple Inc.
+
 from dataclasses import dataclass
-from typing import 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_additive_causal_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -44,7 +45,7 @@ class Attention(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         B, L, D = x.shape
 
@@ -59,8 +60,8 @@ class Attention(nn.Module):
         if cache is not None:
             keys, values = cache.update_and_fetch(keys, values)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
         )
 
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
@@ -98,7 +99,7 @@ class TransformerBlock(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         r = self.attn(self.ln_1(x), mask, cache)
         h = x + r
@@ -124,22 +125,22 @@ class GPT2Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
         _, L = inputs.shape
 
         hidden_states = self.wte(inputs)
 
-        mask = None
-        if hidden_states.shape[1] > 1:
+        offset = 0
+        if cache is not None and len(cache) > 0 and cache[0] is not None:
+            offset = cache[0].offset
 
-            position_ids = mx.array(np.arange(L))
-            hidden_states += self.wpe(position_ids)
+        position_ids = mx.arange(offset, offset + L)
+        hidden_states += self.wpe(position_ids)
 
-            mask = create_additive_causal_mask(
-                hidden_states.shape[1], cache[0].offset if cache is not None else 0
-            )
-            mask = mask.astype(hidden_states.dtype)
+        if mask is None:
+            mask = create_attention_mask(hidden_states, cache)
 
         if cache is None:
             cache = [None] * len(self.h)
@@ -160,9 +161,10 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, cache)
+        out = self.model(inputs, mask, cache)
         out = self.model.wte.as_linear(out)
         return out
 
@@ -197,11 +199,3 @@ class Model(nn.Module):
     @property
     def layers(self):
         return self.model.h
-
-    @property
-    def head_dim(self):
-        return self.args.n_embd // self.args.n_head
-
-    @property
-    def n_kv_heads(self):
-        return self.args.num_key_value_heads

mlx_lm/models/gpt_bigcode.py

@@ -1,11 +1,13 @@
+# Copyright © 2023-2024 Apple Inc.
+
 from dataclasses import dataclass
-from typing import 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_additive_causal_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -55,7 +57,7 @@ class Attention(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         B, L, D = x.shape
 
@@ -72,8 +74,8 @@ class Attention(nn.Module):
         if cache is not None:
             keys, values = cache.update_and_fetch(keys, values)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
         )
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
         return self.c_proj(output)
@@ -112,7 +114,7 @@ class TransformerBlock(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         r = self.attn(self.ln_1(x), mask, cache)
         h = x + r
@@ -135,6 +137,7 @@ class GPTBigCodeModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
         B, L = inputs.shape
@@ -142,18 +145,16 @@ class GPTBigCodeModel(nn.Module):
         hidden_states = self.wte(inputs)
 
         mask = None
-        if hidden_states.shape[1] > 1:
-
-            position_ids = mx.array(np.arange(L))
-            hidden_states += self.wpe(position_ids)
-
-            mask = create_additive_causal_mask(
-                hidden_states.shape[1], cache[0].offset if cache is not None else 0
-            )
-            mask = mask.astype(hidden_states.dtype)
+        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))
+
+        hidden_states += self.wpe(position_ids)
 
         for layer, c in zip(self.h, cache):
             hidden_states = layer(hidden_states, mask, cache=c)
@@ -173,9 +174,10 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
-        out = self.transformer(inputs, cache)
+        out = self.transformer(inputs, mask, cache)
         if self.args.tie_word_embeddings:
             out = self.transformer.wte.as_linear(out)
         else:
@@ -185,11 +187,3 @@ class Model(nn.Module):
     @property
     def layers(self):
         return self.transformer.h
-
-    @property
-    def head_dim(self):
-        return self.args.n_embd // self.args.n_head
-
-    @property
-    def n_kv_heads(self):
-        return self.args.num_key_value_heads

mlx_lm/models/gpt_neox.py

@@ -0,0 +1,218 @@
+# 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
+
+# Based on the transformers implementation at:
+# https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neox/modeling_gpt_neox.py
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    max_position_embeddings: int
+    hidden_size: int
+    num_attention_heads: int
+    num_hidden_layers: int
+    layer_norm_eps: float
+    vocab_size: int
+    rotary_emb_base: int
+    rotary_pct: float
+    num_key_value_heads: int = None
+
+    def __post_init__(self):
+        if self.num_key_value_heads is None:
+            self.num_key_value_heads = self.num_attention_heads
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        assert (
+            args.hidden_size % args.num_attention_heads == 0
+        ), "hidden_size must be divisible by num_attention_heads"
+
+        self.hidden_size = args.hidden_size
+        self.num_attention_heads = args.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_attention_heads
+
+        self.rope = nn.RoPE(
+            dims=int(self.head_dim * args.rotary_pct),
+            traditional=False,
+            base=args.rotary_emb_base,
+        )
+
+        self.scale = self.head_dim**-0.5
+
+        self.query_key_value = nn.Linear(
+            self.hidden_size, 3 * self.hidden_size, bias=True
+        )
+        self.dense = nn.Linear(self.hidden_size, self.hidden_size, bias=True)
+
+    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)
+
+        new_qkv_shape = qkv.shape[:-1] + (self.num_attention_heads, 3 * self.head_dim)
+        qkv = qkv.reshape(*new_qkv_shape)
+
+        queries, keys, values = [x.transpose(0, 2, 1, 3) for x in qkv.split(3, -1)]
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.dense(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        self.hidden_size = args.hidden_size
+        self.dense_h_to_4h = nn.Linear(self.hidden_size, 4 * self.hidden_size)
+        self.dense_4h_to_h = nn.Linear(4 * self.hidden_size, self.hidden_size)
+
+    def __call__(self, x) -> mx.array:
+        # gelu_approx corresponds to FastGELUActivation in transformers.
+        return self.dense_4h_to_h(nn.gelu_approx(self.dense_h_to_4h(x)))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        self.hidden_size = args.hidden_size
+        self.layer_norm_eps = args.layer_norm_eps
+        self.attention = Attention(args)
+        self.mlp = MLP(args)
+        self.input_layernorm = nn.LayerNorm(
+            self.hidden_size,
+            eps=self.layer_norm_eps,
+        )
+        self.post_attention_layernorm = nn.LayerNorm(
+            self.hidden_size, eps=self.layer_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        residual = x
+        # NeoX runs attention and feedforward network in parallel.
+        attn = self.attention(self.input_layernorm(x), mask, cache)
+        ffn = self.mlp(self.post_attention_layernorm(x))
+        out = attn + ffn + residual
+        return out
+
+
+class GPTNeoXModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.hidden_size = args.hidden_size
+        self.vocab_size = args.vocab_size
+        self.num_hidden_layers = args.num_hidden_layers
+        self.layer_norm_eps = args.layer_norm_eps
+        assert self.vocab_size > 0
+        self.embed_in = nn.Embedding(self.vocab_size, self.hidden_size)
+        self.embed_out = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+        self.h = [TransformerBlock(args=args) for _ in range(self.num_hidden_layers)]
+        self.final_layer_norm = nn.LayerNorm(self.hidden_size, eps=self.layer_norm_eps)
+
+    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)
+
+        for layer, c in zip(self.h, cache):
+            hidden_states = layer(hidden_states, mask, cache=c)
+
+        out = self.final_layer_norm(hidden_states)
+        out = self.embed_out(out)
+
+        return out
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = GPTNeoXModel(args)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        out = self.model(inputs, mask, cache)
+        return out
+
+    def sanitize(self, weights):
+        new_weights = {}
+
+        for w_key, w_value in weights.items():
+            # Created through register_buffer in Pytorch, not needed here.
+            ignore_suffixes = [
+                ".attention.bias",
+                ".attention.masked_bias",
+                ".attention.rotary_emb.inv_freq",
+            ]
+
+            skip_weight = False
+            for ignored_suffix in ignore_suffixes:
+                if w_key.endswith(ignored_suffix):
+                    skip_weight = True
+                    break
+
+            if skip_weight:
+                continue
+
+            if not w_key.startswith("model."):
+                w_key = f"model.{w_key}"
+
+            w_key = w_key.replace(".gpt_neox.layers.", ".h.")
+            w_key = w_key.replace(".gpt_neox.", ".")
+
+            new_weights[w_key] = w_value
+
+        return new_weights
+
+    @property
+    def layers(self):
+        return self.model.h

mlx_lm/models/granite.py

@@ -0,0 +1,195 @@
+# Copyright © 2023-2024 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    rms_norm_eps: float
+    vocab_size: int
+    logits_scaling: float
+    attention_multiplier: float
+    embedding_multiplier: float
+    residual_multiplier: float
+    max_position_embeddings: int
+    num_key_value_heads: int
+    attention_bias: bool
+    mlp_bias: bool
+    rope_theta: float
+    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.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)
+
+        # Prepare the queries, keys and values for the attention computation
+        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, 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 = 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 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.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.mlp(self.post_attention_layernorm(h))
+        out = h + r * self.residual_multiplier
+        return out
+
+
+class GraniteModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.num_hidden_layers = args.num_hidden_layers
+        assert self.vocab_size > 0
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            TransformerBlock(args=args) for _ in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.embedding_multiplier = args.embedding_multiplier
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        h = self.embed_tokens(inputs) * self.embedding_multiplier
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, cache=c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = GraniteModel(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+        self.logits_scaling = args.logits_scaling
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        out = self.model(inputs, mask, cache)
+        if self.args.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
+        return out / self.logits_scaling
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/helium.py

@@ -0,0 +1,185 @@
+# 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):
+    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
+    attention_bias: bool
+    head_dim: int
+    max_position_embeddings: int
+    mlp_bias: bool
+    model_type: str
+    rope_theta: float
+    tie_word_embeddings: bool
+
+
+class HeliumAttention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        assert args.num_key_value_heads is not None
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+
+        head_dim = args.hidden_size // n_heads
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=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.rope = nn.RoPE(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)
+
+        # Prepare the queries, keys and values for the attention computation
+        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class HeliumMLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.hidden_size = args.hidden_size
+        self.intermediate_size = args.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) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class HeliumDecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.hidden_size = args.hidden_size
+
+        self.self_attn = HeliumAttention(args)
+        self.mlp = HeliumMLP(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 HeliumModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_hidden_layers = args.num_hidden_layers
+        self.vocab_size = args.vocab_size
+
+        assert self.vocab_size > 0
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+
+        self.layers = [HeliumDecoderLayer(args) for _ in range(args.num_hidden_layers)]
+
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ) -> mx.array:
+        h = self.embed_tokens(inputs)
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+
+        self.model = HeliumModel(args)
+
+        self.vocab_size = args.vocab_size
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ) -> mx.array:
+        out = self.model(inputs, mask, cache)
+        if self.args.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
+        return out
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/hunyuan.py

@@ -0,0 +1,337 @@
+# Copyright © 2023-2024 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Tuple, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .switch_layers import SwitchGLU
+
+
+@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
+    attention_bias: bool
+    moe_topk: int
+    num_experts: int
+    num_shared_expert: int
+    use_mixed_mlp_moe: bool
+    use_qk_norm: bool
+    rms_norm_eps: float
+    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
+
+    def __post_init__(self):
+
+        if self.rope_scaling:
+            required_keys = {"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}")
+
+
+def _int_or_list(arg, idx):
+    if isinstance(arg, list):
+        return arg[idx]
+    return arg
+
+
+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, kv_proj: bool, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        assert args.num_key_value_heads is not None
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+
+        head_dim = args.hidden_size // n_heads
+        self.scale = head_dim**-0.5
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=args.attention_bias)
+        if kv_proj:
+            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)
+
+        self.rope = DynamicNTKAlphaRoPE(
+            head_dim,
+            base=args.rope_theta,
+            scaling_alpha=args.rope_scaling["alpha"],
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+        kv_states=None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries = self.q_proj(x)
+        if kv_states is None:
+            keys, values = self.k_proj(x), self.v_proj(x)
+            kv_states = keys, values
+        else:
+            keys, values = kv_states
+
+        # Prepare the queries, keys and values for the attention computation
+        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        offset = cache.offset if cache else 0
+        queries = self.rope(queries, offset=offset)
+        keys = self.rope(keys, offset=offset)
+        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), kv_states
+
+
+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 Gate(nn.Module):
+    def __init__(self, dim, num_experts):
+        super().__init__()
+        self.wg = nn.Linear(dim, num_experts, bias=False)
+
+    def __call__(self, x) -> mx.array:
+        return self.wg(x)
+
+
+class MoeBlock(nn.Module):
+    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:
+            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 = _int_or_list(args.moe_topk, layer_idx)
+
+        self.gate = Gate(dim, 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,
+        x: mx.array,
+    ):
+        gates = self.gate(x)
+        gates = mx.softmax(gates, axis=-1, precise=True)
+
+        k = self.top_k
+        inds = mx.stop_gradient(mx.argpartition(-gates, kth=k - 1, axis=-1)[..., :k])
+        scores = mx.take_along_axis(gates, inds, axis=-1)
+
+        y = self.switch_mlp(x, inds)
+        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)
+            y = y + shared_expert_output
+
+        return y
+
+
+class DecoderLayer(nn.Module):
+    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, layer_idx)
+
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+        self.args = args
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+        shared_kv_states: Optional[Tuple[mx.array, mx.array]] = None,
+    ):
+        r, shared_kv_states = self.self_attn(
+            self.input_layernorm(x), mask, cache, shared_kv_states
+        )
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        out = h + r
+        return out, shared_kv_states
+
+
+class HunYuanModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.num_hidden_layers = args.num_hidden_layers
+        assert self.vocab_size > 0
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            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)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        h = self.embed_tokens(inputs)
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for 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
+            h, shared_kv_states = layer(h, mask, c, shared_kv_states)
+
+        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 = HunYuanModel(args)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        out = self.model(inputs, mask, cache)
+        return self.model.embed_tokens.as_linear(out)
+
+    def sanitize(self, weights):
+
+        if "model.layers.0.mlp.gate_and_up_proj.weight" in weights:
+            new_weights = {}
+            D = self.args.hidden_size
+            n_kv_heads = self.args.num_key_value_heads
+            n_kv_groups = self.args.num_attention_heads // n_kv_heads
+            head_dim = D // self.args.num_attention_heads
+            for k, v in weights.items():
+                if "qkv_proj" in k:
+                    v = v.reshape(n_kv_heads, n_kv_groups + 2, head_dim, -1)
+                    splits = v.split([n_kv_groups, n_kv_groups + 1], axis=1)
+                    for k_up, v_new in zip(["q_proj", "k_proj", "v_proj"], splits):
+                        k_new = k.replace("qkv_proj", k_up)
+                        new_weights[k_new] = mx.flatten(v_new, 0, 2)
+                elif "gate_and_up_proj" in k:
+                    splits = v.split(2, axis=0)
+                    for k_up, v_new in zip(["up_proj", "gate_proj"], splits):
+                        k_new = k.replace("gate_and_up_proj", k_up)
+                        new_weights[k_new] = v_new
+                else:
+                    new_weights[k] = v
+            weights = new_weights
+
+        if "model.layers.0.mlp.experts.0.up_proj.weight" not in weights:
+            return weights
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}"
+            for n in ["up_proj", "down_proj", "gate_proj"]:
+                for k in ["weight", "scales", "biases"]:
+                    if f"{prefix}.mlp.experts.0.{n}.{k}" in weights:
+                        to_join = [
+                            weights.pop(f"{prefix}.mlp.experts.{e}.{n}.{k}")
+                            for e in range(self.args.num_experts)
+                        ]
+                        weights[f"{prefix}.mlp.switch_mlp.{n}.{k}"] = mx.stack(to_join)
+        return weights
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/internlm2.py

@@ -1,10 +1,12 @@
+# Copyright © 2023-2024 Apple Inc.
+
 from dataclasses import dataclass
-from typing import 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
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -17,6 +19,7 @@ class ModelArgs(BaseModelArgs):
     rms_norm_eps: float
     vocab_size: int
     bias: bool = True
+    max_position_embeddings: int = 32768
     num_key_value_heads: int = None
     rope_theta: float = 10000
     rope_traditional: bool = False
@@ -32,8 +35,50 @@ class ModelArgs(BaseModelArgs):
             if not all(key in self.rope_scaling for key in required_keys):
                 raise ValueError(f"rope_scaling must contain keys {required_keys}")
 
-            if self.rope_scaling["type"] != "linear":
-                raise ValueError("rope_scaling 'type' currently only supports 'linear'")
+            if self.rope_scaling["type"] not in ["linear", "dynamic"]:
+                raise ValueError(
+                    "rope_scaling 'type' currently only supports 'linear' or 'dynamic"
+                )
+
+
+class DynamicNTKScalingRoPE(nn.Module):
+    """Implements the rotary positional encoding with Dynamic NTK scaling."""
+
+    def __init__(
+        self,
+        dims: int,
+        max_position_embeddings: int = 2048,
+        traditional: bool = False,
+        base: float = 10000,
+        scale: float = 1.0,
+    ):
+        super().__init__()
+        self.max_position_embeddings = max_position_embeddings
+        self.original_base = base
+        self.dims = dims
+        self.traditional = traditional
+        self.scale = scale
+
+    def extra_repr(self):
+        return f"{self.dims}, traditional={self.traditional}, max_position_embeddings={self.max_position_embeddings}, scaling_factor={self.scaling_factor}"
+
+    def __call__(self, x, offset: int = 0):
+        seq_len = x.shape[1] + offset
+        if seq_len > self.max_position_embeddings:
+            base = self.original_base * (
+                (self.scale * seq_len / self.max_position_embeddings) - (self.scale - 1)
+            ) ** (self.dims / (self.dims - 2))
+        else:
+            base = self.original_base
+
+        return mx.fast.rope(
+            x,
+            self.dims,
+            traditional=self.traditional,
+            base=base,
+            scale=self.scale,
+            offset=offset,
+        )
 
 
 class Attention(nn.Module):
@@ -56,10 +101,12 @@ class Attention(nn.Module):
         rope_scale = (
             1 / args.rope_scaling["factor"]
             if args.rope_scaling is not None and args.rope_scaling["type"] == "linear"
-            else 1
+            else 2.0
         )
-        self.rope = nn.RoPE(
+
+        self.rope = DynamicNTKScalingRoPE(
             head_dim,
+            max_position_embeddings=args.max_position_embeddings,
             traditional=args.rope_traditional,
             base=args.rope_theta,
             scale=rope_scale,
@@ -69,7 +116,7 @@ class Attention(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         B, L, D = x.shape
 
@@ -94,8 +141,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
         )
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
         return self.wo(output)
@@ -124,7 +171,7 @@ class TransformerBlock(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         r = self.attention(self.attention_norm(x), mask, cache)
         h = x + r
@@ -146,14 +193,13 @@ class InternLM2Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
         h = self.tok_embeddings(inputs)
 
-        mask = None
-        if h.shape[1] > 1:
-            mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
-            mask = mask.astype(h.dtype)
+        if mask is None:
+            mask = create_attention_mask(h, cache)
 
         if cache is None:
             cache = [None] * len(self.layers)
@@ -176,23 +222,20 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, cache)
+        out = self.model(inputs, mask, cache)
         if self.args.tie_word_embeddings:
             out = self.model.tok_embeddings.as_linear(out)
         else:
             out = self.output(out)
         return out
 
+    def sanitize(self, weights):
+        # Remove unused precomputed rotary freqs
+        return {k: v for k, v in weights.items() if "attention.rope.inv_freq" not in k}
+
     @property
     def layers(self):
         return self.model.layers
-
-    @property
-    def head_dim(self):
-        return self.args.hidden_size // self.args.num_attention_heads
-
-    @property
-    def n_kv_heads(self):
-        return self.args.num_key_value_heads

mlx_lm/models/internlm3.py

@@ -0,0 +1,241 @@
+# 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
+
+
+@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
+    bias: bool = False
+    qkv_bias: bool = False
+    max_position_embeddings: int = 32768
+    num_key_value_heads: int = None
+    rope_theta: float = 10000
+    rope_traditional: bool = False
+    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.rope_scaling:
+            required_keys = {"factor", "rope_type"}
+            if not all(key in self.rope_scaling for key in required_keys):
+                raise ValueError(f"rope_scaling must contain keys {required_keys}")
+
+            if self.rope_scaling["rope_type"] not in ["linear", "dynamic"]:
+                raise ValueError(
+                    "rope_scaling 'rope_type' currently only supports 'linear' or 'dynamic"
+                )
+
+
+class DynamicNTKScalingRoPE(nn.Module):
+    """Implements the rotary positional encoding with Dynamic NTK scaling."""
+
+    def __init__(
+        self,
+        dims: int,
+        max_position_embeddings: int = 2048,
+        traditional: bool = False,
+        base: float = 10000,
+        scale: float = 1.0,
+    ):
+        super().__init__()
+        self.max_position_embeddings = max_position_embeddings
+        self.original_base = base
+        self.dims = dims
+        self.traditional = traditional
+        self.scale = scale
+
+    def extra_repr(self):
+        return f"{self.dims}, traditional={self.traditional}, max_position_embeddings={self.max_position_embeddings}, scaling_factor={self.scaling_factor}"
+
+    def __call__(self, x, offset: int = 0):
+        seq_len = x.shape[1] + offset
+        if seq_len > self.max_position_embeddings:
+            base = self.original_base * (
+                (self.scale * seq_len / self.max_position_embeddings) - (self.scale - 1)
+            ) ** (self.dims / (self.dims - 2))
+        else:
+            base = self.original_base
+
+        return mx.fast.rope(
+            x,
+            self.dims,
+            traditional=self.traditional,
+            base=base,
+            scale=self.scale,
+            offset=offset,
+        )
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        qkv_bias = args.qkv_bias
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+        self.n_kv_groups = n_heads // args.num_key_value_heads
+
+        self.head_dim = head_dim = args.hidden_size // n_heads
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=qkv_bias)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=qkv_bias)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=qkv_bias)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=qkv_bias)
+
+        rope_scale = (
+            1 / args.rope_scaling["factor"]
+            if args.rope_scaling is not None
+            and args.rope_scaling["rope_type"] == "linear"
+            else 2.0
+        )
+
+        self.rope = DynamicNTKScalingRoPE(
+            head_dim,
+            max_position_embeddings=args.max_position_embeddings,
+            traditional=args.rope_traditional,
+            base=args.rope_theta,
+            scale=rope_scale,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        # Prepare the queries, keys and values for the attention computation
+        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, dim, hidden_dim, bias):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=bias)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=bias)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=bias)
+
+    def __call__(self, x) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.self_attn = Attention(args)
+        self.mlp = MLP(args.hidden_size, args.intermediate_size, args.bias)
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        out = h + r
+        return out
+
+
+class InternLM2Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        assert args.vocab_size > 0
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            TransformerBlock(args=args) for _ in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        h = self.embed_tokens(inputs)
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, 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 = InternLM2Model(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        out = self.model(inputs, mask, cache)
+        if self.args.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
+        return out
+
+    def sanitize(self, weights):
+        # Remove unused precomputed rotary freqs
+        return {k: v for k, v in weights.items() if "attention.rope.inv_freq" not in k}
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/kimi_vl.py

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

mlx_lm/models/llama.py

@@ -1,10 +1,13 @@
+# Copyright © 2023-2024 Apple Inc.
+
 from dataclasses import dataclass
-from typing import 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, KVCache, create_additive_causal_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
 
 
 @dataclass
@@ -16,6 +19,8 @@ class ModelArgs(BaseModelArgs):
     num_attention_heads: int
     rms_norm_eps: float
     vocab_size: int
+    head_dim: Optional[int] = None
+    max_position_embeddings: Optional[int] = None
     num_key_value_heads: Optional[int] = None
     attention_bias: bool = False
     mlp_bias: bool = False
@@ -28,14 +33,6 @@ class ModelArgs(BaseModelArgs):
         if self.num_key_value_heads is None:
             self.num_key_value_heads = self.num_attention_heads
 
-        if self.rope_scaling:
-            required_keys = {"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}")
-
-            if self.rope_scaling["type"] != "linear":
-                raise ValueError("rope_scaling 'type' currently only supports 'linear'")
-
 
 class Attention(nn.Module):
     def __init__(self, args: ModelArgs):
@@ -45,7 +42,8 @@ class Attention(nn.Module):
         self.n_heads = n_heads = args.num_attention_heads
         self.n_kv_heads = n_kv_heads = args.num_key_value_heads
 
-        head_dim = args.hidden_size // n_heads
+        self.head_dim = head_dim = args.head_dim or args.hidden_size // n_heads
+
         self.scale = head_dim**-0.5
         if hasattr(args, "attention_bias"):
             attention_bias = args.attention_bias
@@ -57,23 +55,19 @@ class Attention(nn.Module):
         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)
 
-        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(
-            head_dim,
-            traditional=args.rope_traditional,
-            base=args.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__(
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[KVCache] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         B, L, D = x.shape
 
@@ -92,9 +86,10 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
         )
+
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
         return self.o_proj(output)
 
@@ -135,7 +130,7 @@ class TransformerBlock(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[KVCache] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         r = self.self_attn(self.input_layernorm(x), mask, cache)
         h = x + r
@@ -160,16 +155,17 @@ class LlamaModel(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)
 
-        mask = None
-        if h.shape[1] > 1:
-            mask = create_additive_causal_mask(
-                h.shape[1], cache[0].offset if cache is not None else 0
-            )
-            mask = mask.astype(h.dtype)
+        if mask is None:
+            mask = create_attention_mask(h, cache)
 
         if cache is None:
             cache = [None] * len(self.layers)
@@ -192,9 +188,11 @@ 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, cache)
+        out = self.model(inputs, mask, cache, input_embeddings)
         if self.args.tie_word_embeddings:
             out = self.model.embed_tokens.as_linear(out)
         else:
@@ -203,18 +201,13 @@ class Model(nn.Module):
 
     def sanitize(self, weights):
         # Remove unused precomputed rotary freqs
-        return {
+        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
-
-    @property
-    def head_dim(self):
-        return self.args.hidden_size // self.args.num_attention_heads
-
-    @property
-    def n_kv_heads(self):
-        return self.args.num_key_value_heads

mlx_lm/models/llama4.py

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

mlx_lm/models/mamba.py

@@ -0,0 +1,242 @@
+# Copyright © 2024-2025 Apple Inc.
+
+import math
+from dataclasses import dataclass
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs
+from .cache import MambaCache
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    state_size: int
+    num_hidden_layers: int
+    conv_kernel: int
+    use_bias: bool
+    use_conv_bias: bool
+    time_step_rank: int
+    tie_word_embeddings: bool = True
+    use_bcdt_rms: bool = False
+    mixer_rms_eps: float = 1e-6
+
+    def __post_init__(self):
+        if not hasattr(self, "hidden_size") and hasattr(self, "d_model"):
+            self.hidden_size = self.d_model
+        if not hasattr(self, "intermediate_size") and hasattr(self, "d_inner"):
+            self.intermediate_size = self.d_inner
+        if not hasattr(self, "state_size") and hasattr(self, "d_state"):
+            self.state_size = self.d_state
+        if not hasattr(self, "num_hidden_layers") and hasattr(self, "n_layer"):
+            self.num_hidden_layers = self.n_layer
+        if not hasattr(self, "num_hidden_layers") and hasattr(self, "n_layers"):
+            self.num_hidden_layers = self.n_layers
+        if not hasattr(self, "conv_kernel") and hasattr(self, "d_conv"):
+            self.conv_kernel = self.d_conv
+        if not hasattr(self, "use_bias") and hasattr(self, "bias"):
+            self.use_bias = self.bias
+        if not hasattr(self, "use_conv_bias") and hasattr(self, "conv_bias"):
+            self.use_conv_bias = self.conv_bias
+
+        if self.time_step_rank == "auto":
+            self.time_step_rank = math.ceil(self.hidden_size / 16)
+        if self.model_type == "falcon_mamba":
+            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__()
+        self.args = args
+
+        self.hidden_size = args.hidden_size
+        self.ssm_state_size = args.state_size
+        self.conv_kernel_size = args.conv_kernel
+        self.intermediate_size = args.intermediate_size
+        self.time_step_rank = int(args.time_step_rank)
+        self.use_conv_bias = args.use_conv_bias
+        self.use_bcdt_rms = args.use_bcdt_rms
+        if self.use_bcdt_rms:
+            self.mixer_norm = lambda x: mx.fast.rms_norm(
+                x, mx.ones(x.shape[-1], x.dtype), eps=args.mixer_rms_eps
+            )
+
+        self.in_proj = nn.Linear(
+            self.hidden_size, self.intermediate_size * 2, bias=args.use_bias
+        )
+
+        self.conv1d = DepthWiseConv1d(
+            channels=self.intermediate_size,
+            kernel_size=self.conv_kernel_size,
+            bias=self.use_conv_bias,
+            padding=self.conv_kernel_size - 1,
+        )
+
+        self.x_proj = nn.Linear(
+            self.intermediate_size,
+            self.time_step_rank + 2 * self.ssm_state_size,
+            bias=False,
+        )
+        self.dt_proj = nn.Linear(self.time_step_rank, self.intermediate_size, bias=True)
+
+        A = mx.repeat(
+            mx.arange(1.0, self.ssm_state_size + 1.0).reshape([1, self.ssm_state_size]),
+            repeats=self.intermediate_size,
+            axis=0,
+        )
+        self.A_log = mx.log(A)
+        self.D = mx.ones([self.intermediate_size])
+
+        self.out_proj = nn.Linear(
+            self.intermediate_size, self.hidden_size, bias=args.use_bias
+        )
+
+    def ssm_step(self, x, A, state=None):
+        D = self.D
+        deltaBC = self.x_proj(x)
+        delta, B, C = map(
+            self.mixer_norm if self.use_bcdt_rms else lambda x: x,
+            mx.split(
+                deltaBC,
+                [self.time_step_rank, self.time_step_rank + self.ssm_state_size],
+                axis=-1,
+            ),
+        )
+        if self.use_bcdt_rms:
+            delta, B, C = map(self.mixer_norm, (delta, B, C))
+        delta = nn.softplus(self.dt_proj(delta))
+        new_state = mx.expand_dims(delta * x, -1) * mx.expand_dims(B, 1)
+        if state is not None:
+            new_state += state * mx.exp(mx.expand_dims(delta, -1) * A)
+        y = (new_state @ mx.expand_dims(C, -1)).squeeze(2)
+        y = y + D * x
+        return y, new_state
+
+    def _process_sequence(self, x, conv_cache, state_cache):
+        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)
+        x = nn.silu(conv_out)
+
+        A = -mx.exp(self.A_log)
+
+        outputs = []
+        current_state = state_cache
+        y = []
+        for t in range(T):
+            y_t, current_state = self.ssm_step(x[:, t], A, current_state)
+            y.append(y_t)
+        y = mx.stack(y, axis=1)
+        z = self.out_proj(nn.silu(z) * y)
+        return z, (new_conv_cache, current_state)
+
+    def __call__(self, x, cache):
+        if cache is None:
+            conv_cache, state_cache = None, None
+        else:
+            conv_cache, state_cache = cache[0], cache[1]
+
+        output, (new_conv_cache, new_state_cache) = self._process_sequence(
+            x, conv_cache, state_cache
+        )
+
+        if isinstance(cache, MambaCache):
+            cache[0] = new_conv_cache
+            cache[1] = new_state_cache
+
+        return output
+
+
+class ResidualBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.mixer = MambaBlock(args)
+        self.norm = nn.RMSNorm(args.hidden_size)
+
+    def __call__(self, x: mx.array, cache):
+        return self.mixer(self.norm(x), cache) + x
+
+
+class Mamba(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embeddings = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [ResidualBlock(args) for _ in range(args.num_hidden_layers)]
+        self.norm_f = nn.RMSNorm(args.hidden_size)
+
+    def __call__(self, x: mx.array, cache):
+        x = self.embeddings(x)
+        if cache is None:
+            cache = [None] * len(self.layers)
+        for layer, c in zip(self.layers, cache):
+            x = layer(x, c)
+        return self.norm_f(x)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.backbone = Mamba(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):
+        B, T = inputs.shape
+
+        x = self.backbone(inputs, cache)
+
+        if self.args.tie_word_embeddings:
+            logits = self.backbone.embeddings.as_linear(x)
+        else:
+            logits = self.lm_head(x)
+
+        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

mlx_lm/models/mimo.py

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

mlx_lm/models/minicpm.py

@@ -1,11 +1,13 @@
+# Copyright © 2023-2025 Apple Inc.
+
 from dataclasses import dataclass
-from typing import Dict, Optional, Tuple, Union
+from typing import Any, Dict, Optional, Union
 
 import mlx.core as mx
 import mlx.nn as nn
-import numpy as np
 
-from .base import BaseModelArgs
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
 
 
 @dataclass
@@ -21,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
@@ -66,24 +69,19 @@ 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__(
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ):
         B, L, _ = x.shape
 
@@ -103,8 +101,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        attn_output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        attn_output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
         )
 
         attn_output = attn_output.transpose(0, 2, 1, 3).reshape(B, L, -1)
@@ -133,12 +131,12 @@ class DecoderLayer(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         r = self.self_attn(self.input_layernorm(x), mask, cache)
-        h = x + r * (self.scale_depth / np.sqrt(self.num_hidden_layers))
+        h = x + r * (self.scale_depth / self.num_hidden_layers**0.5)
         r = self.mlp(self.post_attention_layernorm(h))
-        out = h + r * (self.scale_depth / np.sqrt(self.num_hidden_layers))
+        out = h + r * (self.scale_depth / self.num_hidden_layers**0.5)
         return out
 
 
@@ -156,14 +154,13 @@ 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
 
-        mask = None
-        if h.shape[1] > 1:
-            mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
-            mask = mask.astype(h.dtype)
+        if mask is None:
+            mask = create_attention_mask(h, cache)
 
         if cache is None:
             cache = [None] * len(self.layers)
@@ -187,9 +184,10 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, cache)
+        out = self.model(inputs, mask, cache)
 
         if not self.args.tie_word_embeddings:
             out = self.lm_head(out / (self.args.hidden_size / self.args.dim_model_base))
@@ -206,11 +204,3 @@ class Model(nn.Module):
     @property
     def layers(self):
         return self.model.layers
-
-    @property
-    def head_dim(self):
-        return self.args.hidden_size // self.args.num_attention_heads
-
-    @property
-    def n_kv_heads(self):
-        return self.args.num_key_value_heads

mlx_lm/models/minicpm3.py

@@ -0,0 +1,250 @@
+# Copyright © 2023-2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import SuScaledRoPE
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    dim_model_base: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    rms_norm_eps: float
+    vocab_size: int
+    num_key_value_heads: int
+    q_lora_rank: int
+    qk_nope_head_dim: int
+    qk_rope_head_dim: int
+    kv_lora_rank: int
+    scale_depth: float
+    scale_emb: float
+    max_position_embeddings: int
+    attention_bias: bool = False
+    rope_theta: float = 1000000.0
+    rope_traditional: bool = False
+    rope_scaling: Optional[Dict[str, Union[str, float]]] = None
+    tie_word_embeddings: bool = False
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+
+        self.qk_rope_head_dim = self.args.qk_rope_head_dim
+        self.qk_nope_head_dim = self.args.qk_nope_head_dim
+        self.attention_bias = self.args.attention_bias
+        self.kv_lora_rank = self.args.kv_lora_rank
+        self.num_heads = self.args.num_attention_heads
+        self.q_lora_rank = self.args.q_lora_rank
+        self.hidden_size = self.args.hidden_size
+
+        self.v_head_dim = self.hidden_size // self.args.num_attention_heads
+        self.q_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+        self.softmax_scale = self.q_head_dim ** (-0.5)
+
+        self.q_a_proj = nn.Linear(
+            self.hidden_size, self.q_lora_rank, bias=self.attention_bias
+        )
+        self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank)
+
+        self.q_b_proj = nn.Linear(
+            self.q_lora_rank, self.num_heads * self.q_head_dim, bias=False
+        )
+
+        self.kv_a_proj_with_mqa = nn.Linear(
+            self.hidden_size,
+            self.kv_lora_rank + self.qk_rope_head_dim,
+            bias=self.attention_bias,
+        )
+
+        self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank)
+
+        self.kv_b_proj = nn.Linear(
+            self.kv_lora_rank,
+            self.num_heads
+            * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
+            bias=False,
+        )
+
+        self.o_proj = nn.Linear(
+            self.num_heads * self.v_head_dim,
+            self.hidden_size,
+            bias=self.attention_bias,
+        )
+
+        self.rope = SuScaledRoPE(
+            dims=args.qk_rope_head_dim,
+            base=args.rope_theta,
+            max_position_embeddings=args.max_position_embeddings,
+            original_max_position_embeddings=args.rope_scaling.get(
+                "original_max_position_embeddings", 4096
+            ),
+            short_factor=args.rope_scaling.get("short_factor", 1.0),
+            long_factor=args.rope_scaling.get("long_factor", 1.0),
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Dict[str, mx.array]] = None,
+    ):
+        B, L, _ = x.shape
+
+        # Project query
+        q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x)))
+        q = q.reshape(B, L, self.num_heads, -1).transpose(0, 2, 1, 3)
+        q_nope, q_pe = mx.split(q, [self.qk_nope_head_dim], axis=-1)
+
+        # Project key and value
+        compressed_kv = self.kv_a_proj_with_mqa(x)
+        compressed_kv, k_pe = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
+        k_pe = k_pe.reshape(B, L, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3)
+
+        kv = self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
+        kv = kv.reshape(B, L, self.num_heads, -1).transpose(0, 2, 1, 3)
+
+        k_nope, values = mx.split(kv, [self.qk_nope_head_dim], axis=-1)
+
+        # Apply RoPE to the query and key parts that need position embedding
+        if cache is not None:
+            q_pe = self.rope(q_pe, offset=cache.offset)
+            k_pe = self.rope(k_pe, offset=cache.offset)
+        else:
+            q_pe = self.rope(q_pe)
+            k_pe = self.rope(k_pe)
+
+        # Create the full query and key tensors by combining the parts
+        # Broadcast k_pe to all heads
+        k_pe_broadcasted = mx.broadcast_to(
+            k_pe, (B, self.num_heads, L, self.qk_rope_head_dim)
+        )
+
+        # Use concatenate for queries
+        queries = mx.concatenate([q_nope, q_pe], axis=-1)
+
+        # Use concatenate for keys
+        keys = mx.concatenate([k_nope, k_pe_broadcasted], axis=-1)
+
+        # Update cache if needed
+        if cache is not None:
+            keys, values = cache.update_and_fetch(keys, values)
+
+        # Perform attention
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.softmax_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):
+        super().__init__()
+        self.gate_proj = nn.Linear(args.hidden_size, args.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(args.hidden_size, args.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(args.intermediate_size, args.hidden_size, bias=False)
+
+    def __call__(self, x):
+        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.args = args
+        self.hidden_size = args.hidden_size
+        self.num_hidden_layers = args.num_hidden_layers
+
+        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.scale_depth = args.scale_depth
+        self.num_hidden_layers = args.num_hidden_layers
+
+    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.scale_depth / (self.num_hidden_layers**0.5))
+        r = self.mlp(self.post_attention_layernorm(h))
+        out = h + r * (self.scale_depth / (self.num_hidden_layers**0.5))
+        return out
+
+
+class MiniCPM3Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        assert self.vocab_size > 0
+
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [DecoderLayer(args) for _ in range(args.num_hidden_layers)]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        h = self.embed_tokens(inputs) * self.args.scale_emb
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = MiniCPM3Model(args)
+
+        if not self.args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        out = self.model(inputs, mask, cache)
+
+        if not self.args.tie_word_embeddings:
+            out = self.lm_head(out / (self.args.hidden_size / self.args.dim_model_base))
+        else:
+            out = self.model.embed_tokens.as_linear(out)
+
+        return out
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/mistral3.py

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

mlx_lm/models/mixtral.py

@@ -1,11 +1,12 @@
-import math
+# Copyright © 2023-2024 Apple Inc.
+
 from dataclasses import dataclass
-from typing import 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
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .switch_layers import SwitchGLU
 
 
@@ -64,7 +65,7 @@ class MixtralAttention(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         B, L, D = x.shape
 
@@ -85,8 +86,8 @@ class MixtralAttention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
         )
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
         return self.o_proj(output)
@@ -136,7 +137,7 @@ class MixtralDecoderLayer(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         r = self.self_attn(self.input_layernorm(x), mask, cache)
         h = x + r
@@ -160,15 +161,13 @@ class MixtralModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
         h = self.embed_tokens(inputs)
 
-        mask = None
-        T = h.shape[1]
-        if T > 1:
-            mask = nn.MultiHeadAttention.create_additive_causal_mask(T)
-            mask = mask.astype(h.dtype)
+        if mask is None:
+            mask = create_attention_mask(h, cache)
 
         if cache is None:
             cache = [None] * len(self.layers)
@@ -190,9 +189,10 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, cache)
+        out = self.model(inputs, mask, cache)
         return self.lm_head(out)
 
     def sanitize(self, weights):
@@ -217,11 +217,3 @@ class Model(nn.Module):
     @property
     def layers(self):
         return self.model.layers
-
-    @property
-    def head_dim(self):
-        return self.args.hidden_size // self.args.num_attention_heads
-
-    @property
-    def n_kv_heads(self):
-        return self.args.num_key_value_heads

mlx_lm/models/nemotron-nas.py

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

mlx_lm/models/nemotron.py

@@ -0,0 +1,220 @@
+# Copyright © 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
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    hidden_act: str
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    norm_eps: float
+    vocab_size: int
+    num_key_value_heads: int
+    head_dim: Optional[int] = None
+    max_position_embeddings: Optional[int] = None
+    attention_bias: bool = False
+    mlp_bias: bool = False
+    partial_rotary_factor: float = 0.5
+    rope_theta: float = 10000.0
+    rope_traditional: bool = False
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = False
+
+    def __post_init__(self):
+        if self.rope_scaling:
+            if not "factor" in self.rope_scaling:
+                raise ValueError(f"rope_scaling must contain 'factor'")
+            rope_type = self.rope_scaling.get("type") or self.rope_scaling.get(
+                "rope_type"
+            )
+            if rope_type is None:
+                raise ValueError(
+                    f"rope_scaling must contain either 'type' or 'rope_type'"
+                )
+            if rope_type not in ["linear"]:
+                raise ValueError("rope_scaling 'type' currently only supports 'linear'")
+
+
+@partial(mx.compile, shapeless=True)
+def relu_squared(x):
+    return nn.relu(x).square()
+
+
+class NemotronLayerNorm1P(nn.LayerNorm):
+    def __call__(self, x):
+        weight = self.weight + 1 if "weight" in self else None
+        bias = self.bias if "bias" in self else None
+        return mx.fast.layer_norm(x, weight, bias, self.eps)
+
+
+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.partial_rotary_factor = args.partial_rotary_factor
+
+        self.scale = head_dim**-0.5
+        if hasattr(args, "attention_bias"):
+            attention_bias = args.attention_bias
+        else:
+            attention_bias = False
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=attention_bias)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=attention_bias)
+
+        rope_scale = 1.0
+        if args.rope_scaling and args.rope_scaling["type"] == "linear":
+            assert isinstance(args.rope_scaling["factor"], float)
+            rope_scale = 1 / args.rope_scaling["factor"]
+        self.rope = nn.RoPE(
+            int(self.partial_rotary_factor * self.head_dim),
+            base=args.rope_theta,
+            scale=rope_scale,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, _ = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        # Prepare the queries, keys and values for the attention computation
+        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        hidden_dim = args.intermediate_size
+        mlp_bias = args.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(relu_squared(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 = NemotronLayerNorm1P(args.hidden_size, eps=args.norm_eps)
+        self.post_attention_layernorm = NemotronLayerNorm1P(
+            args.hidden_size, eps=args.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 NemotronModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.num_hidden_layers = args.num_hidden_layers
+        assert self.vocab_size > 0
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            TransformerBlock(args=args) for _ in range(args.num_hidden_layers)
+        ]
+        self.norm = NemotronLayerNorm1P(args.hidden_size, eps=args.norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        h = self.embed_tokens(inputs)
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, cache=c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = NemotronModel(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        out = self.model(inputs, mask, cache)
+        if self.args.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
+        return out
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/olmo.py

@@ -1,17 +1,19 @@
+# Copyright © 2023-2024 Apple Inc.
+
+import sys
 from dataclasses import dataclass
-from sys import exit
-from typing import Optional, Tuple
+from typing import Any, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs
+from .base import BaseModelArgs, create_attention_mask
 
 try:
     import hf_olmo
 except ImportError:
     print("To run olmo install ai2-olmo: pip install ai2-olmo")
-    exit(1)
+    sys.exit(1)
 
 
 @dataclass
@@ -66,7 +68,7 @@ class TransformerBlock(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         B, L, D = x.shape
 
@@ -96,7 +98,7 @@ class TransformerBlock(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         r = self.attend(self.att_norm(x), mask, cache)
         h = x + r
@@ -122,14 +124,13 @@ class Transformer(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
         h = self.wte(inputs)
 
-        mask = None
-        if h.shape[1] > 1:
-            mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
-            mask = mask.astype(h.dtype)
+        if mask is None:
+            mask = create_attention_mask(h, cache)
 
         if cache is None:
             cache = [None] * len(self.blocks)
@@ -153,9 +154,10 @@ class OlmoModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
-        return self.transformer(inputs, cache)
+        return self.transformer(inputs, mask, cache)
 
 
 class Model(nn.Module):
@@ -168,18 +170,11 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
-        return self.model(inputs, cache)
+        return self.model(inputs, mask, cache)
 
     @property
     def layers(self):
         return self.model.transformer.blocks
-
-    @property
-    def head_dim(self):
-        return self.args.d_model // self.args.n_heads
-
-    @property
-    def n_kv_heads(self):
-        return self.args.n_heads

mlx_lm/models/olmo2.py

@@ -0,0 +1,212 @@
+# Copyright © 2023-2024 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    rms_norm_eps: float
+    vocab_size: int
+    head_dim: Optional[int] = None
+    max_position_embeddings: Optional[int] = None
+    num_key_value_heads: 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
+
+    def __post_init__(self):
+        if self.num_key_value_heads is None:
+            self.num_key_value_heads = self.num_attention_heads
+
+
+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
+        if hasattr(args, "attention_bias"):
+            attention_bias = args.attention_bias
+        else:
+            attention_bias = False
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=attention_bias)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=attention_bias)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=attention_bias)
+
+        self.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            args.rope_traditional,
+            args.rope_scaling,
+            args.max_position_embeddings,
+        )
+
+        self.q_norm = nn.RMSNorm(n_heads * head_dim, args.rms_norm_eps)
+        self.k_norm = nn.RMSNorm(n_kv_heads * head_dim, 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 = self.q_norm(queries)
+        keys = self.k_norm(keys)
+
+        # Prepare the queries, keys and values for the attention computation
+        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, 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 = 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 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.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 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
+        assert self.vocab_size > 0
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            TransformerBlock(args=args) for _ in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        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)
+
+        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,
+        mask=None,
+    ):
+        out = self.model(inputs, cache, mask)
+        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
+        return {
+            k: v for k, v in weights.items() if "self_attn.rotary_emb.inv_freq" not in k
+        }
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/olmoe.py

@@ -0,0 +1,217 @@
+# 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
+    num_experts: int
+    num_experts_per_tok: int
+    norm_topk_prob: bool = False
+    head_dim: Optional[int] = None
+    max_position_embeddings: Optional[int] = None
+    num_key_value_heads: 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
+
+    def __post_init__(self):
+        if self.num_key_value_heads is None:
+            self.num_key_value_heads = self.num_attention_heads
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+
+        self.head_dim = head_dim = args.head_dim or args.hidden_size // n_heads
+
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=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.rope = initialize_rope(
+            self.head_dim,
+            args.rope_theta,
+            args.rope_traditional,
+            args.rope_scaling,
+            args.max_position_embeddings,
+        )
+
+        self.q_norm = nn.RMSNorm(n_heads * head_dim, args.rms_norm_eps)
+        self.k_norm = nn.RMSNorm(n_kv_heads * head_dim, 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 = self.q_norm(queries)
+        keys = self.k_norm(keys)
+        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 OlmoeSparseMoeBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_experts = args.num_experts
+        self.top_k = args.num_experts_per_tok
+        self.norm_topk_prob = args.norm_topk_prob
+
+        self.gate = nn.Linear(args.hidden_size, self.num_experts, bias=False)
+        self.switch_mlp = SwitchGLU(
+            args.hidden_size,
+            args.intermediate_size,
+            self.num_experts,
+            bias=args.mlp_bias,
+        )
+
+    def __call__(self, x: mx.array) -> mx.array:
+        B, L, D = x.shape
+        x_flat = x.reshape(-1, D)
+        router_logits = self.gate(x_flat)
+        routing_weights = mx.softmax(router_logits, axis=1, precise=True)
+        k = self.top_k
+        indices = mx.stop_gradient(
+            mx.argpartition(-routing_weights, kth=k - 1, axis=-1)[..., :k]
+        )
+        scores = mx.take_along_axis(routing_weights, indices, axis=-1)
+        if self.norm_topk_prob:
+            scores = scores / scores.sum(axis=-1, keepdims=True)
+        y = self.switch_mlp(x_flat, indices)
+        y = (y * scores[..., None]).sum(axis=-2)
+        return y.reshape(B, L, D)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.self_attn = Attention(args)
+        self.mlp = OlmoeSparseMoeBlock(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:
+        x = x + self.self_attn(self.input_layernorm(x), mask, cache)
+        x = x + self.mlp(self.post_attention_layernorm(x))
+        return x
+
+
+class OlmoeModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.vocab_size = args.vocab_size
+        self.num_hidden_layers = args.num_hidden_layers
+        assert self.vocab_size > 0
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            TransformerBlock(args=args) for _ in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        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)
+        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 = OlmoeModel(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,
+        mask=None,
+    ):
+        out = self.model(inputs, cache, mask)
+        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 "model.layers.0.mlp.experts.0.up_proj.weight" not in weights:
+            return weights
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}"
+            for n in ["up_proj", "down_proj", "gate_proj"]:
+                for k in ["weight", "scales", "biases"]:
+                    if f"{prefix}.mlp.experts.0.{n}.{k}" in weights:
+                        to_join = [
+                            weights.pop(f"{prefix}.mlp.experts.{e}.{n}.{k}")
+                            for e in range(self.args.num_experts)
+                        ]
+                        weights[f"{prefix}.mlp.switch_mlp.{n}.{k}"] = mx.stack(to_join)
+        return weights
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/openelm.py

@@ -1,10 +1,12 @@
+# Copyright © 2023-2024 Apple Inc.
+
 from dataclasses import dataclass
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Union
 
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -78,7 +80,7 @@ class Attention(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         B, L, D = x.shape
 
@@ -105,8 +107,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
         )
 
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
@@ -150,7 +152,7 @@ class TransformerBlock(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         r = self.attn(self.attn_norm(x), mask, cache)
         h = x + r
@@ -176,14 +178,13 @@ class OpenELMModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
         h = self.token_embeddings(inputs)
 
-        mask = None
-        if h.shape[1] > 1:
-            mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
-            mask = mask.astype(h.dtype)
+        if mask is None:
+            mask = create_attention_mask(h, cache)
 
         if cache is None:
             cache = [None] * len(self.layers)
@@ -206,9 +207,10 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
-        out = self.transformer(inputs, cache)
+        out = self.transformer(inputs, mask, cache)
         if self.args.share_input_output_layers:
             out = self.transformer.token_embeddings.as_linear(out)
         else:
@@ -219,11 +221,3 @@ class Model(nn.Module):
     @property
     def layers(self):
         return self.transformer.layers
-
-    @property
-    def head_dim(self):
-        return self.args.head_dim
-
-    @property
-    def n_kv_heads(self):
-        return self.args.num_kv_heads

mlx_lm/models/openlm.py

@@ -1,182 +0,0 @@
-from dataclasses import dataclass
-from typing import Optional, Tuple
-
-import mlx.core as mx
-import mlx.nn as nn
-
-from .base import BaseModelArgs, create_additive_causal_mask
-
-
-@dataclass
-class ParamsArgs(BaseModelArgs):
-    dim: int
-    ffn_type: str
-    n_heads: int
-    n_layers: int
-    norm_eps: float
-    positional_embedding_type: str
-    post_embed_norm: bool
-    qk_norm: bool
-    vocab_size: int
-    weight_tying: bool
-
-
-@dataclass
-class ModelArgs(BaseModelArgs):
-    model_type: str
-    params_args_dict: ParamsArgs
-
-
-class Attention(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-
-        self.dim = args.dim
-        self.n_heads = args.n_heads
-        self.head_dim = self.dim // self.n_heads
-        self.qk_norm = args.qk_norm
-        self.scale = self.head_dim**-0.5
-
-        self.in_proj = nn.Linear(self.dim, 3 * self.dim, bias=False)
-        self.out_proj = nn.Linear(self.dim, self.dim, bias=False)
-        if self.qk_norm:
-            self.q_norm = nn.LayerNorm(args.dim, eps=args.norm_eps, bias=False)
-            self.k_norm = nn.LayerNorm(args.dim, eps=args.norm_eps, bias=False)
-        self.rope = nn.RoPE(
-            self.head_dim,
-            traditional=False,
-        )
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
-    ) -> mx.array:
-        B, L, D = x.shape
-
-        queries, keys, values = self.in_proj(x).split(3, axis=-1)
-
-        if self.qk_norm:
-            queries = self.q_norm(queries)
-            keys = self.q_norm(keys)
-
-        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
-        keys = keys.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
-        values = values.reshape(B, L, self.n_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 = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
-        )
-
-        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
-        return self.out_proj(output)
-
-
-class MLP(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        # https://github.com/mlfoundations/open_lm/blob/c65b43042ff31c0fe26f930decf1ccab1b03ab4b/open_lm/model.py#L254C2-L254C3
-        hidden_dim = 256 * ((int(2 * 4 * args.dim / 3) + 256 - 1) // 256)
-        self.w12 = nn.Linear(args.dim, 2 * hidden_dim, bias=False)
-        self.w3 = nn.Linear(hidden_dim, args.dim, bias=False)
-
-    def __call__(self, x) -> mx.array:
-        gate, x = self.w12(x).split(2, axis=-1)
-        return self.w3(nn.silu(gate) * x)
-
-
-class TransformerBlock(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.attention = Attention(args)
-        self.feed_forward = MLP(args)
-        self.ffn_norm = nn.LayerNorm(args.dim, eps=args.norm_eps, bias=False)
-        self.attention_norm = nn.LayerNorm(args.dim, eps=args.norm_eps, bias=False)
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
-    ) -> mx.array:
-        r = self.attention(self.attention_norm(x), mask, cache)
-        h = x + r
-        r = self.feed_forward(self.ffn_norm(h))
-        out = h + r
-        return out
-
-
-class OpenLM(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.tok_embeddings = nn.Embedding(args.vocab_size, args.dim)
-        self.layers = [TransformerBlock(args=args) for _ in range(args.n_layers)]
-        self.norm = nn.LayerNorm(args.dim, eps=args.norm_eps, bias=False)
-        self.output = nn.Linear(args.dim, args.vocab_size, bias=False)
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        cache=None,
-    ):
-        _, L = inputs.shape
-
-        h = self.tok_embeddings(inputs)
-
-        mask = None
-        if h.shape[1] > 1:
-            mask = create_additive_causal_mask(
-                h.shape[1], cache[0].offset if cache is not None else 0
-            )
-            mask = mask.astype(h.dtype)
-
-        if cache is None:
-            cache = [None] * len(self.layers)
-
-        for layer, c in zip(self.layers, cache):
-            h = layer(h, mask, cache=c)
-
-        return self.output(self.norm(h))
-
-
-class Model(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        args.params_args_dict = ParamsArgs.from_dict(args.params_args_dict)
-        self.args = args.params_args_dict
-        self.model_type = args.model_type
-        self.model = OpenLM(self.args)
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        cache=None,
-    ):
-        out = self.model(inputs, cache)
-        return out
-
-    def sanitize(self, weights):
-        # Remove unused precomputed rotary freqs
-        return {k: v for k, v in weights.items() if "inv_freq" not in k}
-
-    @property
-    def layers(self):
-        return self.model.layers
-
-    @property
-    def head_dim(self):
-        return self.args.dim // self.args.n_heads
-
-    @property
-    def n_kv_heads(self):
-        return self.args.n_heads

mlx_lm/models/phi.py

@@ -1,11 +1,12 @@
+# Copyright © 2023-2024 Apple Inc.
+
 import math
 from dataclasses import dataclass
-from typing import Tuple
 
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -91,8 +92,13 @@ class PhiAttention(nn.Module):
             keys = self.rope(keys)
 
         scale = math.sqrt(1 / queries.shape[-1])
-        output = mx.fast.scaled_dot_product_attention(
-            queries.astype(mx.float32), keys, values, scale=scale, mask=mask
+        output = scaled_dot_product_attention(
+            queries.astype(mx.float32),
+            keys,
+            values,
+            cache=cache,
+            scale=scale,
+            mask=mask,
         ).astype(values.dtype)
 
         output = output.moveaxis(2, 1).reshape(B, L, -1)
@@ -105,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):
@@ -136,16 +141,15 @@ class PhiModel(nn.Module):
             config.hidden_size, eps=config.layer_norm_eps
         )
 
-    def __call__(self, x, cache):
+    def __call__(self, x, mask, 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 = None
-        if x.shape[1] > 1:
-            mask = nn.MultiHeadAttention.create_additive_causal_mask(x.shape[1])
-            mask = mask.astype(x.dtype)
-
         for layer, c in zip(self.layers, cache):
             x = layer(x, mask, c)
         return self.final_layernorm(x)
@@ -162,19 +166,12 @@ class Model(nn.Module):
     def __call__(
         self,
         x: mx.array,
-        cache: mx.array = None,
-    ) -> Tuple[mx.array, mx.array]:
-        y = self.model(x, cache)
+        mask: mx.array = None,
+        cache=None,
+    ) -> mx.array:
+        y = self.model(x, mask, cache)
         return self.lm_head(y)
 
     @property
     def layers(self):
         return self.model.layers
-
-    @property
-    def head_dim(self):
-        return self.args.hidden_size // self.args.num_attention_heads
-
-    @property
-    def n_kv_heads(self):
-        return self.args.num_key_value_heads

mlx_lm/models/phi3.py

@@ -1,11 +1,13 @@
+# Copyright © 2023-2024 Apple Inc.
+
 from dataclasses import dataclass
-from typing import Dict, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs
-from .su_rope import SuScaledRotaryEmbedding
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import SuScaledRoPE
 
 
 @dataclass
@@ -17,12 +19,14 @@ class ModelArgs(BaseModelArgs):
     num_attention_heads: int
     rms_norm_eps: float
     vocab_size: int
-    num_key_value_heads: int = None
+    num_key_value_heads: Optional[int] = None
     rope_theta: float = 10000
     rope_traditional: bool = False
-    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    rope_scaling: Optional[Dict[str, Union[float, List[float]]]] = None
+    partial_rotary_factor: float = 1.0
     max_position_embeddings: int = 131072
     original_max_position_embeddings: int = 4096
+    tie_word_embeddings: bool = False
 
     def __post_init__(self):
         if self.num_key_value_heads is None:
@@ -33,9 +37,9 @@ class ModelArgs(BaseModelArgs):
             if not all(key in self.rope_scaling for key in required_keys):
                 raise ValueError(f"rope_scaling must contain keys {required_keys}")
 
-            if self.rope_scaling["type"] not in ["su", "linear"]:
+            if self.rope_scaling["type"] not in ["longrope", "su", "linear"]:
                 print(
-                    "[WARNING] rope_scaling 'type' currently only supports 'linear' and 'su'; setting rope scaling to false."
+                    "[WARNING] rope_scaling 'type' currently only supports 'linear', 'su', and 'longrope'; setting rope scaling to false."
                 )
                 self.rope_scaling = None
 
@@ -46,6 +50,7 @@ class Attention(nn.Module):
 
         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
         self.num_hidden_layers = args.num_hidden_layers
 
@@ -56,23 +61,23 @@ class Attention(nn.Module):
         self.qkv_proj = nn.Linear(dim, op_size, bias=False)
         self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
 
-        rope_scale = 1.0
-        if args.rope_scaling and args.rope_scaling["type"] == "su":
-            self.rope = SuScaledRotaryEmbedding(
-                head_dim,
-                traditional=False,
+        rope_dim = int(head_dim * args.partial_rotary_factor)
+        if args.rope_scaling and args.rope_scaling["type"] in ["longrope", "su"]:
+            self.rope = SuScaledRoPE(
+                rope_dim,
                 base=args.rope_theta,
-                scale=rope_scale,
                 max_position_embeddings=args.max_position_embeddings,
                 original_max_position_embeddings=args.original_max_position_embeddings,
                 short_factor=args.rope_scaling["short_factor"],
                 long_factor=args.rope_scaling["long_factor"],
             )
         else:
+            rope_scale = 1.0
             if args.rope_scaling and args.rope_scaling["type"] == "linear":
+                assert isinstance(args.rope_scaling["factor"], float)
                 rope_scale = 1 / args.rope_scaling["factor"]
             self.rope = nn.RoPE(
-                head_dim,
+                rope_dim,
                 traditional=args.rope_traditional,
                 base=args.rope_theta,
                 scale=rope_scale,
@@ -82,7 +87,7 @@ class Attention(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         B, L, D = x.shape
 
@@ -105,8 +110,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
         )
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
         return self.o_proj(output)
@@ -141,7 +146,7 @@ class TransformerBlock(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         r = self.self_attn(self.input_layernorm(x), mask, cache)
         h = x + r
@@ -166,14 +171,13 @@ class Phi3Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
         h = self.embed_tokens(inputs)
 
-        mask = None
-        if h.shape[1] > 1:
-            mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
-            mask = mask.astype(h.dtype)
+        if mask is None:
+            mask = create_attention_mask(h, cache)
 
         if cache is None:
             cache = [None] * len(self.layers)
@@ -189,25 +193,23 @@ class Model(nn.Module):
         super().__init__()
         self.model_type = args.model_type
         self.model = Phi3Model(args)
-        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
         self.args = args
 
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, cache)
-        return self.lm_head(out)
+        out = self.model(inputs, mask, cache)
+        if self.args.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
+        return out
 
     @property
     def layers(self):
         return self.model.layers
-
-    @property
-    def head_dim(self):
-        return self.args.hidden_size // self.args.num_attention_heads
-
-    @property
-    def n_kv_heads(self):
-        return self.args.num_key_value_heads

mlx_lm/models/phi3small.py

@@ -1,11 +1,14 @@
+# Copyright © 2023-2024 Apple Inc.
+
+import math
 from dataclasses import dataclass
 from functools import partial
-from typing import Dict, Optional, Tuple, Union
+from typing import Any, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -19,14 +22,14 @@ class ModelArgs(BaseModelArgs):
     num_attention_heads: int
     layer_norm_epsilon: float
     vocab_size: int
-    num_key_value_heads: int = None
+    num_key_value_heads: int
     mup_attn_multiplier: float = 1.0
     mup_use_scaling: bool = True
     mup_embedding_multiplier: float = 10.0
     mup_width_multiplier: float = 8.0
     rope_embedding_base: float = 1000000
     rope_position_scale: float = 1.0
-    blocksparse_block_size: int = (64,)
+    blocksparse_block_size: int = 64
     blocksparse_num_local_blocks: int = 16
     blocksparse_vert_stride: int = 8
 
@@ -157,7 +160,7 @@ class Attention(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         B, L, D = x.shape
 
@@ -185,8 +188,8 @@ class Attention(nn.Module):
                 queries, keys, values, scale=self.scale, mask=mask
             )
         else:
-            output = mx.fast.scaled_dot_product_attention(
-                queries, keys, values, scale=self.scale, mask=mask
+            output = scaled_dot_product_attention(
+                queries, keys, values, cache=cache, scale=self.scale, mask=mask
             )
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
         return self.dense(output)
@@ -226,7 +229,7 @@ class TransformerBlock(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         r = self.self_attn(self.input_layernorm(x), mask, cache)
         h = x + r
@@ -255,16 +258,15 @@ 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
 
-        mask = None
-        if h.shape[1] > 1:
-            mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
-            mask = mask.astype(h.dtype)
+        if mask is None:
+            mask = create_attention_mask(h, cache, return_array=True)
 
         if cache is None:
             cache = [None] * len(self.layers)
@@ -290,9 +292,10 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, cache)
+        out = self.model(inputs, mask, cache)
         out = self.model.embed_tokens.as_linear(out)
         if self.mup_width_multiplier:
             out = out / self.mup_width_multiplier
@@ -303,16 +306,8 @@ class Model(nn.Module):
     def layers(self):
         return self.model.layers
 
-    @property
-    def head_dim(self):
-        return self.args.hidden_size // self.args.num_attention_heads
-
     def sanitize(self, weights):
         # Remove unused precomputed rotary freqs
         return {
             k: v for k, v in weights.items() if "self_attn.rotary_emb.inv_freq" not in k
         }
-
-    @property
-    def n_kv_heads(self):
-        return self.args.num_key_value_heads

mlx_lm/models/phimoe.py

@@ -0,0 +1,214 @@
+# Copyright © 2024 Apple Inc.
+import math
+from dataclasses import dataclass
+from typing import 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 SuScaledRoPE
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str = "phimoe"
+    vocab_size: int = 32064
+    hidden_size: int = 4096
+    intermediate_size: int = 6400
+    num_hidden_layers: int = 32
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 8
+    max_position_embeddings: int = 131072
+    original_max_position_embeddings: int = 4096
+    rms_norm_eps: float = 1e-6
+    rope_scaling: Dict[str, Union[float, List[float]]] = None
+    num_local_experts: int = 16
+    num_experts_per_tok: int = 2
+    rope_theta: float = 10000.0
+
+
+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.hidden_size // n_heads
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=True)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=True)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=True)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=True)
+
+        self.rope = SuScaledRoPE(
+            head_dim,
+            base=args.rope_theta,
+            max_position_embeddings=args.max_position_embeddings,
+            original_max_position_embeddings=args.original_max_position_embeddings,
+            short_factor=args.rope_scaling["short_factor"],
+            long_factor=args.rope_scaling["long_factor"],
+            short_mscale=args.rope_scaling["short_mscale"],
+            long_mscale=args.rope_scaling["long_mscale"],
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache=None,
+    ) -> mx.array:
+        B, L, D = x.shape
+
+        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
+
+        # Prepare the queries, keys and values for the attention computation
+        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class PhiMoESparseMoeBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.hidden_dim = args.hidden_size
+        self.ffn_dim = args.intermediate_size
+        self.num_experts = args.num_local_experts
+        self.top_k = args.num_experts_per_tok
+
+        self.gate = nn.Linear(self.hidden_dim, self.num_experts, bias=False)
+        self.switch_mlp = SwitchGLU(self.hidden_dim, self.ffn_dim, self.num_experts)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        gates = self.gate(x)
+
+        k = self.top_k
+        inds = mx.stop_gradient(mx.argpartition(-gates, kth=k - 1, axis=-1)[..., :k])
+        scores = mx.take_along_axis(gates, inds, axis=-1)
+        scores = mx.softmax(scores, axis=-1, precise=True)
+
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2)
+
+        return y
+
+
+class PhiMoEDecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.hidden_size = args.hidden_size
+
+        self.self_attn = Attention(args)
+        self.block_sparse_moe = PhiMoESparseMoeBlock(args)
+        self.input_layernorm = nn.LayerNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.LayerNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache=None,
+    ) -> mx.array:
+        residual = x
+        hidden_states = self.input_layernorm(x)
+        hidden_states = self.self_attn(hidden_states, mask=mask, cache=cache)
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.block_sparse_moe(hidden_states)
+        hidden_states = residual + hidden_states
+
+        return hidden_states
+
+
+class PhiMoEModel(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 = [PhiMoEDecoderLayer(args) for _ in range(args.num_hidden_layers)]
+        self.norm = nn.LayerNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ) -> mx.array:
+        h = self.embed_tokens(inputs)
+
+        if mask is None:
+            mask = create_attention_mask(h, cache)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.model_type = args.model_type
+        self.args = args
+        self.model = PhiMoEModel(args)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=True)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        out = self.model(inputs, mask, cache)
+        return self.lm_head(out)
+
+    def sanitize(self, weights):
+        if "model.layers.0.block_sparse_moe.experts.0.w1.weight" not in weights:
+            return weights
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"model.layers.{l}"
+            for n, m in [("w1", "gate_proj"), ("w2", "down_proj"), ("w3", "up_proj")]:
+                for k in ["weight", "scales", "biases"]:
+                    if f"{prefix}.block_sparse_moe.experts.0.{n}.{k}" in weights:
+                        to_join = [
+                            weights.pop(
+                                f"{prefix}.block_sparse_moe.experts.{e}.{n}.{k}"
+                            )
+                            for e in range(self.args.num_local_experts)
+                        ]
+                        weights[f"{prefix}.block_sparse_moe.switch_mlp.{m}.{k}"] = (
+                            mx.stack(to_join)
+                        )
+
+        return weights
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/phixtral.py

@@ -1,3 +1,5 @@
+# Copyright © 2023-2024 Apple Inc.
+
 import inspect
 import math
 from dataclasses import dataclass
@@ -6,6 +8,7 @@ from typing import Tuple
 import mlx.core as mx
 import mlx.nn as nn
 
+from .base import create_attention_mask, scaled_dot_product_attention
 from .switch_layers import SwitchMLP
 
 
@@ -68,8 +71,13 @@ class RoPEAttention(nn.Module):
         # Finally perform the attention computation
         scale = math.sqrt(1 / queries.shape[-1])
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries.astype(mx.float32), keys, values, scale=scale, mask=mask
+        output = scaled_dot_product_attention(
+            queries.astype(mx.float32),
+            keys,
+            values,
+            cache=cache,
+            scale=scale,
+            mask=mask,
         ).astype(values.dtype)
         output = output.moveaxis(2, 1).reshape(B, L, -1)
 
@@ -165,12 +173,11 @@ class Model(nn.Module):
         self,
         x: mx.array,
         mask: mx.array = None,
-        cache: mx.array = None,
-    ) -> Tuple[mx.array, mx.array]:
-        mask = None
-        if x.shape[1] > 1:
-            mask = nn.MultiHeadAttention.create_additive_causal_mask(x.shape[1])
-            mask = mask.astype(x.dtype)
+        cache=None,
+    ) -> mx.array:
+
+        if mask is None:
+            mask = create_attention_mask(x, cache)
 
         y = self.transformer(x, mask, cache)
         return self.lm_head(y)
@@ -193,11 +200,3 @@ class Model(nn.Module):
     @property
     def layers(self):
         return self.transformer.h
-
-    @property
-    def head_dim(self):
-        return self.args.model_dim // self.args.num_heads
-
-    @property
-    def n_kv_heads(self):
-        return self.args.num_heads

mlx_lm/models/pixtral.py

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

mlx_lm/models/plamo.py

@@ -1,11 +1,13 @@
+# Copyright © 2023-2024 Apple Inc.
+
 from dataclasses import dataclass
-from typing import Any, List, 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
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -60,8 +62,8 @@ class Attention(nn.Module):
         self,
         hidden_states: mx.array,
         attention_mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
-    ) -> Tuple[mx.array, Tuple[mx.array, mx.array]]:
+        cache: Optional[Any] = None,
+    ) -> mx.array:
         bsz, q_len, _ = hidden_states.shape
 
         queries = self.q_proj(hidden_states)
@@ -87,10 +89,14 @@ class Attention(nn.Module):
             queries = self.rotary_emb(queries)
             keys = self.rotary_emb(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
+        keys = mx.tile(keys, [1, self.config.n_shared_head, 1, 1])
+        values = mx.tile(values, [1, self.config.n_shared_head, 1, 1])
+
+        output = scaled_dot_product_attention(
             queries,
             keys,
             values,
+            cache=cache,
             scale=self.scale,
             mask=attention_mask,
         )
@@ -125,8 +131,8 @@ class PlamoDecoderLayer(nn.Module):
         self,
         hidden_states: mx.array,
         attention_mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
-    ) -> Tuple[Any, ...]:
+        cache: Optional[Any] = None,
+    ):
         # from LlamaDecoder
         residual = hidden_states
 
@@ -167,14 +173,13 @@ class PlamoModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        cache: Optional[List[Union[Tuple[mx.array, mx.array], None]]] = None,
-    ) -> Tuple[mx.array, Optional[List[Union[Tuple[mx.array, mx.array], None]]]]:
+        cache: Optional[Any] = None,
+        mask: Optional[mx.array] = None,
+    ) -> mx.array:
         h = self.embed_tokens(inputs)
 
-        mask = None
-        if h.shape[1] > 1:
-            mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
-            mask = mask.astype(self.embed_tokens.weight.dtype)
+        if mask is None:
+            mask = create_attention_mask(h, cache)
 
         if cache is None:
             cache = [None for _ in range(len(self.layers.layers))]
@@ -198,19 +203,12 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
-        cache: Optional[List[Tuple[mx.array, mx.array]]] = None,
-    ) -> Tuple[mx.array, mx.array]:
-        out = self.model(inputs, cache)
+        cache: Optional[Any] = None,
+        mask: Optional[mx.array] = None,
+    ) -> mx.array:
+        out = self.model(inputs, cache, mask)
         return self.lm_head(out)
 
     @property
     def layers(self):
         return self.model.layers.layers
-
-    @property
-    def head_dim(self):
-        return self.args.hidden_size // self.args.num_attention_heads
-
-    @property
-    def n_kv_heads(self):
-        return self.args.num_attention_heads // self.args.n_shared_head

mlx_lm/models/plamo2.py

@@ -0,0 +1,599 @@
+# Copyright © 2025 Apple Inc.
+
+import math
+from dataclasses import dataclass
+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 .cache import KVCache, MambaCache
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str = "plamo2"
+    hidden_size: int = 4096
+    num_hidden_layers: int = 32
+    rms_norm_eps: float = 1e-6
+    tie_word_embeddings: bool = True
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 4
+    hidden_size_per_head: int = 128
+    max_position_embeddings: int = 2048
+    attention_window_size: int = 2048
+    full_attention_idx: Optional[list[int]] = None
+    mamba_d_state: int = 64
+    mamba_d_conv: int = 4
+    mamba_num_heads: int = 64
+    mamba_step: int = 2
+    mamba_chunk_size: int = 256
+    mamba_enabled: bool = True
+    intermediate_size: int = 13312
+    vocab_size: int = 32000
+
+
+class RMSNorm(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        eps: float = 1e-6,
+        offset: float = 1.0,
+    ) -> None:
+        super().__init__()
+        self.weight = mx.zeros(hidden_size)
+        self.variance_epsilon = eps
+        self.offset = offset
+
+    def __call__(self, hidden_states: mx.array) -> mx.array:
+        return mx.fast.rms_norm(
+            hidden_states, self.weight + self.offset, self.variance_epsilon
+        )
+
+
+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
+    state_len = conv_state.shape[-2]
+    x = mx.concatenate([conv_state, x], axis=-2)
+    conv_state = x[:, -state_len:]
+    out = mx.conv1d(
+        x,
+        weight,
+        padding=0,
+        groups=dim,
+    )[:, -seqlen:]
+    return nn.silu(out), conv_state
+
+
+class Mamba(nn.Module):
+    def __init__(self, config: ModelArgs) -> None:
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.d_state = config.mamba_d_state
+        self.d_conv = config.mamba_d_conv
+        self.chunk_size = config.mamba_chunk_size
+        self.num_heads = config.mamba_num_heads
+        self.hidden_size_per_head = config.hidden_size_per_head
+
+        self.intermediate_size = self.num_heads * self.hidden_size_per_head
+
+        self.in_proj = nn.Linear(
+            self.hidden_size, 2 * self.intermediate_size, bias=False
+        )
+        self.conv1d = nn.Conv1d(
+            in_channels=self.intermediate_size,
+            out_channels=self.intermediate_size,
+            bias=False,
+            kernel_size=self.d_conv,
+            groups=self.intermediate_size,
+            padding=0,
+        )
+        self.dt_dim = max(64, self.hidden_size // 16)
+        self.bcdt_proj = nn.Linear(
+            self.intermediate_size,
+            self.dt_dim + 2 * self.d_state,
+            bias=False,
+        )
+        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_norm_weight = mx.ones(self.dt_dim)
+        self.B_norm_weight = mx.ones(self.d_state)
+        self.C_norm_weight = mx.ones(self.d_state)
+
+        self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+
+    def __call__(
+        self,
+        hidden_states: mx.array,
+        mask: Optional[mx.array] = None,
+        cache=None,
+    ):
+        bsize, length, _ = hidden_states.shape
+
+        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)
+        # z: (bsize, length, num_heads, hidden_size_per_head)
+        # x: (bsize, length, num_heads, hidden_size_per_head)
+        z, x = mx.split(
+            zx,
+            [
+                self.hidden_size_per_head,
+            ],
+            axis=-1,
+        )
+
+        x = x.reshape(bsize, -1, self.num_heads * self.hidden_size_per_head)
+        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,)
+        dt = mx.fast.rms_norm(dt, self.dt_norm_weight, self.config.rms_norm_eps)
+        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(
+            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,
+        )
+
+        if cache is not None:
+            cache[0] = conv_state
+            cache[1] = ssm_state
+        y = self.out_proj(out.reshape(bsize, length, -1))
+
+        return y
+
+
+class Attention(nn.Module):
+    def __init__(self, config: ModelArgs) -> None:
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        head_dim = config.hidden_size_per_head
+        self.max_position_embeddings = config.max_position_embeddings
+        self.scale = head_dim**-0.5
+
+        self.q_num_heads = config.num_attention_heads
+        self.qk_dim = self.v_dim = head_dim
+        self.k_num_heads = self.v_num_heads = config.num_key_value_heads
+        assert self.q_num_heads % self.k_num_heads == 0
+        self.n_group = self.q_num_heads // self.k_num_heads
+
+        self.q_proj_dim = self.q_num_heads * self.qk_dim
+        self.k_proj_dim = self.k_num_heads * self.qk_dim
+        self.v_proj_dim = self.k_num_heads * self.v_dim
+        self.qkv_proj = nn.Linear(
+            self.hidden_size,
+            self.q_proj_dim + self.k_proj_dim + self.v_proj_dim,
+            bias=False,
+        )
+        self.o_proj = nn.Linear(
+            self.q_num_heads * self.v_dim, self.hidden_size, bias=False
+        )
+
+        self.q_weight = mx.ones((self.q_num_heads, self.qk_dim))
+        self.k_weight = mx.ones((self.k_num_heads, self.qk_dim))
+
+        self.rope = nn.RoPE(self.qk_dim)
+
+    def __call__(
+        self,
+        hidden_states: mx.array,
+        mask: Optional[mx.array] = None,
+        cache=None,
+    ):
+        B, T, _ = hidden_states.shape
+
+        qkv = self.qkv_proj(hidden_states)
+        q, k, v = mx.split(
+            qkv, [self.q_proj_dim, self.q_proj_dim + self.k_proj_dim], axis=-1
+        )
+        q = q.reshape(B, T, self.q_num_heads, self.qk_dim).transpose(0, 2, 1, 3)
+        k = k.reshape(B, T, self.k_num_heads, self.qk_dim).transpose(0, 2, 1, 3)
+        v = v.reshape(B, T, self.v_num_heads, self.v_dim).transpose(0, 2, 1, 3)
+
+        q = mx.fast.rms_norm(q, weight=None, eps=1e-6) * self.q_weight[:, None]
+        k = mx.fast.rms_norm(k, weight=None, eps=1e-6) * self.k_weight[:, None]
+
+        if cache is not None:
+            q = self.rope(q, offset=cache.offset)
+            k = self.rope(k, offset=cache.offset)
+            k, v = cache.update_and_fetch(k, v)
+        else:
+            q = self.rope(q)
+            k = self.rope(k)
+
+        output = mx.fast.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            scale=self.scale,
+            mask=mask,
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(
+            B, T, self.q_num_heads * self.v_dim
+        )
+        return self.o_proj(output)
+
+
+class MLP(nn.Module):
+    def __init__(self, config: ModelArgs) -> None:
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.intermediate_size = config.intermediate_size
+        self.gate_up_proj = nn.Linear(
+            self.hidden_size, self.intermediate_size * 2, bias=False
+        )
+        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        h = self.gate_up_proj(x)
+        hs = mx.split(h, 2, axis=-1)
+        return self.down_proj(nn.silu(hs[0]) * hs[1])
+
+
+class PlamoDecoderLayer(nn.Module):
+    def __init__(self, config: ModelArgs, is_mamba: bool) -> None:
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.is_mamba = is_mamba
+        self.mixer: nn.Module
+        if is_mamba:
+            self.mixer = Mamba(config)
+        else:
+            self.mixer = Attention(config)
+        self.mlp = MLP(config)
+        self.pre_mixer_norm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps, offset=1.0
+        )
+        self.post_mixer_norm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps, offset=1.0 / 5
+        )
+        self.pre_mlp_norm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps, offset=1.0
+        )
+        self.post_mlp_norm = RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps, offset=1.0 / (5**1.5)
+        )
+
+    def __call__(
+        self,
+        hidden_states: mx.array,
+        mask: Optional[mx.array] = None,
+        cache=None,
+    ):
+        residual = hidden_states
+        hidden_states = self.pre_mixer_norm(hidden_states)
+
+        hidden_states_sa = self.mixer(
+            hidden_states=hidden_states,
+            mask=mask,
+            cache=cache,
+        )
+
+        hidden_states_sa = self.post_mixer_norm(hidden_states_sa)
+        hidden_states = residual + hidden_states_sa
+
+        residual = hidden_states
+        hidden_states = self.pre_mlp_norm(hidden_states)
+
+        # Fully Connected
+        hidden_states_mlp = self.mlp(hidden_states)
+
+        # Residual
+        hidden_states_mlp = self.post_mlp_norm(hidden_states_mlp)
+        return residual + hidden_states_mlp
+
+
+def is_mamba(config: ModelArgs, i: int) -> bool:
+    if not config.mamba_enabled:
+        return False
+    assert config.mamba_step > 1
+    assert i < config.num_hidden_layers
+
+    if config.num_hidden_layers <= (config.mamba_step // 2):
+        # use attention in last layer
+        return i != config.num_hidden_layers - 1
+    return (i % config.mamba_step) != (config.mamba_step // 2)
+
+
+class PlamoDecoder(nn.Module):
+    def __init__(self, config: ModelArgs) -> None:
+        super().__init__()
+
+        self.layers = [
+            PlamoDecoderLayer(config, is_mamba=is_mamba(config, i))
+            for i in range(config.num_hidden_layers)
+        ]
+
+    def __call__(self, x: mx.array, mask: mx.array, cache):
+        for i, decoder_layer in enumerate(self.layers):
+            x = decoder_layer(
+                x,
+                mask=mask,
+                cache=cache[i],
+            )
+        return x
+
+
+class PlamoModel(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+
+        self.config = config
+        self.vocab_size = config.vocab_size
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = PlamoDecoder(config)  # type: ignore
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    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,
+        )
+
+        return self.norm(out)
+
+
+class Model(nn.Module):
+    def __init__(self, config: ModelArgs) -> None:
+        super().__init__()
+        self.config = config
+        self.model_type = config.model_type
+        self.model = PlamoModel(config)
+
+        self.vocab_size = config.vocab_size
+
+        if not config.tie_word_embeddings:
+            self.lm_head: nn.Module = nn.Linear(
+                config.hidden_size, self.vocab_size, bias=False
+            )
+
+    def sanitize(self, weights: dict[Any, Any]) -> dict[Any, Any]:
+        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):
+        # TODO use RotatingKVCache is not full_attn
+        # 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:
+        outputs = self.model(
+            inputs=inputs,
+            mask=None,
+            cache=cache,
+        )
+        if self.config.tie_word_embeddings:
+            logits = self.model.embed_tokens.as_linear(outputs)
+        else:
+            logits = self.lm_head(outputs)
+
+        return logits
+
+    @property
+    def layers(self):
+        return self.model.layers.layers

mlx_lm/models/qwen.py

@@ -1,10 +1,11 @@
+# Copyright © 2023-2024 Apple Inc.
+
 from dataclasses import dataclass
-from typing import Tuple
 
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -63,8 +64,8 @@ class Attention(nn.Module):
             queries = self.rotary_emb(queries)
             keys = self.rotary_emb(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
         )
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
 
@@ -122,11 +123,8 @@ class QwenModel(nn.Module):
     def __call__(self, inputs, mask=None, cache=None):
         x = self.wte(inputs)
 
-        mask = None
-        T = x.shape[1]
-        if T > 1:
-            mask = nn.MultiHeadAttention.create_additive_causal_mask(T)
-            mask = mask.astype(x.dtype)
+        if mask is None:
+            mask = create_attention_mask(x, cache)
 
         if cache is None:
             cache = [None] * len(self.h)
@@ -151,19 +149,11 @@ class Model(nn.Module):
         self,
         x: mx.array,
         mask: mx.array = None,
-        cache: mx.array = None,
-    ) -> Tuple[mx.array, mx.array]:
+        cache=None,
+    ) -> mx.array:
         y = self.transformer(x, mask, cache)
         return self.lm_head(y)
 
     @property
     def layers(self):
         return self.transformer.h
-
-    @property
-    def head_dim(self):
-        return self.args.hidden_size // self.args.num_attention_heads
-
-    @property
-    def n_kv_heads(self):
-        return self.args.num_attention_heads

mlx_lm/models/qwen2.py

@@ -1,10 +1,13 @@
+# Copyright © 2023-2024 Apple Inc.
+
 from dataclasses import dataclass
-from typing import 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
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .rope_utils import initialize_rope
 
 
 @dataclass
@@ -16,24 +19,13 @@ class ModelArgs(BaseModelArgs):
     num_attention_heads: int
     rms_norm_eps: float
     vocab_size: int
-    num_key_value_heads: int = None
+    num_key_value_heads: int
+    max_position_embeddings: int = 32768
     rope_theta: float = 1000000
     rope_traditional: bool = False
     rope_scaling: Optional[Dict[str, Union[float, str]]] = None
     tie_word_embeddings: bool = True
 
-    def __post_init__(self):
-        if self.num_key_value_heads is None:
-            self.num_key_value_heads = self.num_attention_heads
-
-        if self.rope_scaling:
-            required_keys = {"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}")
-
-            if self.rope_scaling["type"] != "linear":
-                raise ValueError("rope_scaling 'type' currently only supports 'linear'")
-
 
 class Attention(nn.Module):
     def __init__(self, args: ModelArgs):
@@ -41,6 +33,7 @@ class Attention(nn.Module):
 
         dim = args.hidden_size
         self.n_heads = n_heads = args.num_attention_heads
+        assert args.num_key_value_heads is not None
         self.n_kv_heads = n_kv_heads = args.num_key_value_heads
 
         head_dim = args.hidden_size // n_heads
@@ -51,23 +44,19 @@ 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=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(
+        self.rope = initialize_rope(
             head_dim,
-            traditional=args.rope_traditional,
             base=args.rope_theta,
-            scale=rope_scale,
+            traditional=args.rope_traditional,
+            scaling_config=args.rope_scaling,
+            max_position_embeddings=args.max_position_embeddings,
         )
 
     def __call__(
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         B, L, D = x.shape
 
@@ -86,8 +75,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
         )
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
         return self.o_proj(output)
@@ -121,7 +110,7 @@ class TransformerBlock(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         r = self.self_attn(self.input_layernorm(x), mask, cache)
         h = x + r
@@ -146,14 +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 input_embeddings is not None:
+            h = input_embeddings
+        else:
+            h = self.embed_tokens(inputs)
 
-        mask = None
-        if h.shape[1] > 1:
-            mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
-            mask = mask.astype(h.dtype)
+        if mask is None:
+            mask = create_attention_mask(h, cache)
 
         if cache is None:
             cache = [None] * len(self.layers)
@@ -176,9 +168,11 @@ 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, cache)
+        out = self.model(inputs, mask, cache, input_embeddings)
         if self.args.tie_word_embeddings:
             out = self.model.embed_tokens.as_linear(out)
         else:
@@ -196,11 +190,3 @@ class Model(nn.Module):
     @property
     def layers(self):
         return self.model.layers
-
-    @property
-    def head_dim(self):
-        return self.args.hidden_size // self.args.num_attention_heads
-
-    @property
-    def n_kv_heads(self):
-        return self.args.num_key_value_heads

mlx_lm/models/qwen2_moe.py

@@ -1,11 +1,12 @@
-import math
+# Copyright © 2023-2024 Apple Inc.
+
 from dataclasses import dataclass
-from typing import 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
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .switch_layers import SwitchGLU
 
 
@@ -22,7 +23,7 @@ class ModelArgs(BaseModelArgs):
     shared_expert_intermediate_size: int
     rms_norm_eps: float
     vocab_size: int
-    num_key_value_heads: int = None
+    num_key_value_heads: Optional[int] = None
     rope_theta: float = 1000000
     rope_traditional: bool = False
     rope_scaling: Optional[Dict[str, Union[float, str]]] = None
@@ -47,6 +48,7 @@ class Attention(nn.Module):
 
         dim = args.hidden_size
         self.n_heads = n_heads = args.num_attention_heads
+        assert args.num_key_value_heads is not None
         self.n_kv_heads = n_kv_heads = args.num_key_value_heads
 
         head_dim = args.hidden_size // n_heads
@@ -67,7 +69,7 @@ class Attention(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         B, L, D = x.shape
 
@@ -86,8 +88,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
         )
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
         return self.o_proj(output)
@@ -159,7 +161,7 @@ class Qwen2MoeDecoderLayer(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         r = self.self_attn(self.input_layernorm(x), mask, cache)
         h = x + r
@@ -184,14 +186,13 @@ class Qwen2MoeModel(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
         h = self.embed_tokens(inputs)
 
-        mask = None
-        if h.shape[1] > 1:
-            mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
-            mask = mask.astype(h.dtype)
+        if mask is None:
+            mask = create_attention_mask(h, cache)
 
         if cache is None:
             cache = [None] * len(self.layers)
@@ -213,9 +214,10 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, cache)
+        out = self.model(inputs, mask, cache)
         return self.lm_head(out)
 
     def sanitize(self, weights):
@@ -236,11 +238,3 @@ class Model(nn.Module):
     @property
     def layers(self):
         return self.model.layers
-
-    @property
-    def head_dim(self):
-        return self.args.hidden_size // self.args.num_attention_heads
-
-    @property
-    def n_kv_heads(self):
-        return self.args.num_key_value_heads

mlx_lm/models/qwen3.py

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

mlx_lm/models/qwen3_moe.py

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

mlx_lm/models/recurrent_gemma.py

@@ -0,0 +1,458 @@
+# Copyright © 2023-2024 Apple Inc.
+
+import math
+from dataclasses import dataclass
+from typing import List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import MambaCache, RotatingKVCache
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    attention_bias: bool
+    conv1d_width: int
+    hidden_size: int
+    intermediate_size: int
+    logits_soft_cap: float
+    num_attention_heads: int
+    num_hidden_layers: int
+    num_key_value_heads: int
+    rms_norm_eps: float
+    rope_theta: float
+    attention_window_size: int
+    vocab_size: int
+    embeddings_scale_by_sqrt_dim: bool = True
+    block_types: Optional[List[str]] = None
+    _block_types: Optional[List[str]] = None
+
+    def __post_init__(self):
+        # For some reason these have different names in 2B and 9B
+        if self.block_types is None:
+            self.block_types = self._block_types
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dims: int, eps: float = 1e-5):
+        super().__init__()
+        self.weight = mx.ones((dims,))
+        self.eps = eps
+
+    def __call__(self, x):
+        return mx.fast.rms_norm(x, 1.0 + self.weight, self.eps)
+
+
+def rnn_scan(x, a, h0):
+    assert x.ndim == 3
+    assert a.shape == x.shape[-a.ndim :]
+    assert a.dtype == x.dtype
+
+    if x.shape[1] == 1:
+        # Using scan in sampling mode.
+        if h0 is None:
+            return x, x[:, 0]
+
+        else:
+            y = a * h0[:, None] + x
+            return y, y[:, -1]
+
+    else:
+        # Using scan in linear mode.
+        if h0 is not None:
+            h_t = h0
+        else:
+            B, _, D = x.shape
+            h_t = mx.zeros((B, D), dtype=x.dtype)
+
+        y = mx.zeros_like(x)
+        for t in range(x.shape[1]):
+            h_t = a[:, t] * h_t + x[:, t]
+            y[:, t] = h_t
+
+    return y, h_t
+
+
+class Conv1d(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        kernel_size: int,
+    ):
+        super().__init__()
+        self.weight = mx.zeros((channels, kernel_size, 1))
+        self.bias = mx.zeros((channels,))
+
+    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)
+        y = y + self.bias
+
+        return y, x[:, -K + 1 :, :]
+
+
+class RGLRU(nn.Module):
+    """A Real-Gated Linear Recurrent Unit (RG-LRU) layer."""
+
+    def __init__(
+        self,
+        width: int,
+        num_heads: int,
+    ):
+        super().__init__()
+        self.width = width
+        self.num_heads = num_heads
+        self.head_dim = self.width // self.num_heads
+
+        self.recurrent_param = mx.zeros((self.width,))
+
+        self.input_gate_weight = mx.zeros(
+            (self.num_heads, self.head_dim, self.head_dim),
+        )
+        self.input_gate_bias = mx.zeros((self.num_heads, self.head_dim))
+
+        self.recurrent_gate_weight = mx.zeros(
+            (self.num_heads, self.head_dim, self.head_dim),
+        )
+        self.recurrent_gate_bias = mx.zeros((self.num_heads, self.head_dim))
+
+    def __call__(
+        self,
+        x: mx.array,
+        cache=None,
+    ):
+        B, L, _ = x.shape
+
+        def apply_block_linear(h, w, b):
+            h = h.reshape((B, L, self.num_heads, self.head_dim))
+            h = (h.swapaxes(1, 2) @ w).swapaxes(1, 2) + b
+            return mx.sigmoid(h.flatten(2, 3))
+
+        # Gates for x and a.
+        gate_x = apply_block_linear(x, self.input_gate_weight, self.input_gate_bias)
+        gate_a = apply_block_linear(
+            x, self.recurrent_gate_weight, self.recurrent_gate_bias
+        )
+
+        # Compute the parameter `A` of the recurrence.
+        log_a = -8.0 * gate_a * nn.softplus(self.recurrent_param)
+        a = mx.exp(log_a)
+        a_square = mx.exp(2 * log_a)
+
+        # Gate the input.
+        gated_x = x * gate_x
+
+        # Apply gamma normalization to the input.
+        multiplier = mx.sqrt(1 - a_square)
+        if cache is None:
+            multiplier[:, 0, :] = 1.0
+        normalized_x = gated_x * multiplier.astype(x.dtype)
+
+        y, last_h = rnn_scan(
+            x=normalized_x,
+            a=a,
+            h0=cache,
+        )
+
+        return y, last_h
+
+
+class RecurrentBlock(nn.Module):
+
+    def __init__(
+        self,
+        width: int,
+        num_heads: int,
+        lru_width: int = None,
+        conv1d_temporal_width: int = 4,
+    ):
+        super().__init__()
+        self.width = width
+        self.num_heads = num_heads
+        self.lru_width = lru_width or width
+        self.conv1d_temporal_width = conv1d_temporal_width
+
+        self.linear_y = nn.Linear(width, self.lru_width)
+        self.linear_x = nn.Linear(width, self.lru_width)
+        self.linear_out = nn.Linear(self.lru_width, width)
+        self.conv_1d = Conv1d(
+            channels=self.lru_width,
+            kernel_size=self.conv1d_temporal_width,
+        )
+        self.rg_lru = RGLRU(
+            width=self.lru_width,
+            num_heads=self.num_heads,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        cache=None,
+        mask=None,
+    ):
+        # y branch.
+        y = self.linear_y(x)
+        y = nn.gelu_approx(y)
+
+        # x branch.
+        x = self.linear_x(x)
+        if cache is None:
+            cache = [None, None]
+        x, cache[0] = self.conv_1d(x=x, cache=cache[0])
+        x, cache[1] = self.rg_lru(x=x, cache=cache[1])
+
+        x = x * y
+        x = self.linear_out(x)
+
+        return x
+
+
+class LocalAttentionBlock(nn.Module):
+
+    def __init__(
+        self,
+        width: int,
+        num_heads: int,
+        window_size: int,
+    ):
+        super().__init__()
+        self.width = width
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.scale = (width // num_heads) ** (-0.5)
+
+        self.head_dim = self.width // self.num_heads
+        self.q_proj = nn.Linear(self.width, self.width, bias=False)
+        self.k_proj = nn.Linear(self.width, self.head_dim, bias=False)
+        self.v_proj = nn.Linear(self.width, self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.width, self.width, bias=True)
+        self.rope = nn.RoPE(
+            self.head_dim // 2,
+            traditional=False,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        cache=None,
+        mask=None,
+    ):
+        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_heads, -1).transpose(0, 2, 1, 3)
+        keys = keys.reshape(B, L, 1, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, 1, -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 MLPBlock(nn.Module):
+
+    def __init__(self, width: int, expanded_width: int):
+        super().__init__()
+        self.up_proj = nn.Linear(width, expanded_width // 2)
+        self.gate_proj = nn.Linear(width, expanded_width // 2)
+        self.down_proj = nn.Linear(expanded_width // 2, width)
+
+    def __call__(self, x: mx.array):
+        gate = self.gate_proj(x)
+        x = self.up_proj(x)
+        return self.down_proj(nn.gelu_approx(gate) * x)
+
+
+class ResidualBlock(nn.Module):
+
+    def __init__(
+        self,
+        width: int,
+        mlp_expanded_width: int,
+        num_heads: int,
+        attention_window_size: int,
+        temporal_block_type: str,
+        lru_width: Optional[int] = None,
+        conv1d_temporal_width: int = 4,
+    ):
+        """Initializes the residual block.
+
+        Args:
+          width: The width of the block.
+          mlp_expanded_width: The width of the expansion inside the MLP block.
+          num_heads: The number of heads for the Attention or the RG-LRU.
+          attention_window_size: The window size for the local attention block.
+          temporal_block_type: Either "recurrent" or "attention", specifying the
+            type of recurrent block to use.
+          lru_width: The width of the RG-LRU if different from `width`.
+          conv1d_temporal_width: The width of the temporal convolution.
+        """
+        super().__init__()
+        self.width = width
+        self.mlp_expanded_width = mlp_expanded_width
+        self.num_heads = num_heads
+        self.attention_window_size = attention_window_size
+        self.temporal_block_type = temporal_block_type
+        self.lru_width = lru_width
+        self.conv1d_temporal_width = conv1d_temporal_width
+
+        self.temporal_pre_norm = RMSNorm(width)
+        if self.temporal_block_type == "recurrent":
+            self.temporal_block = RecurrentBlock(
+                width=self.width,
+                num_heads=self.num_heads,
+                lru_width=self.lru_width,
+                conv1d_temporal_width=self.conv1d_temporal_width,
+            )
+
+        else:
+            self.temporal_block = LocalAttentionBlock(
+                width=self.width,
+                num_heads=self.num_heads,
+                window_size=self.attention_window_size,
+            )
+
+        self.channel_pre_norm = RMSNorm(width)
+        self.mlp_block = MLPBlock(
+            width=self.width,
+            expanded_width=self.mlp_expanded_width,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        cache=None,
+        mask=None,
+    ):
+        raw_x = x
+
+        inputs_normalized = self.temporal_pre_norm(raw_x)
+
+        x = self.temporal_block(inputs_normalized, cache=cache, mask=mask)
+        residual = x + raw_x
+
+        x = self.channel_pre_norm(residual)
+        x = self.mlp_block(x)
+
+        x = x + residual
+
+        return x
+
+
+class Griffin(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+
+        self.config = config
+        self.embed_tokens = nn.Embedding(
+            config.vocab_size,
+            config.hidden_size,
+        )
+
+        self.scale_by_sqrt_dim = config.embeddings_scale_by_sqrt_dim
+        block_types = config.block_types
+
+        self.layers = [
+            ResidualBlock(
+                width=config.hidden_size,
+                mlp_expanded_width=config.intermediate_size,
+                num_heads=config.num_attention_heads,
+                attention_window_size=config.attention_window_size,
+                temporal_block_type=block_types[i % len(block_types)],
+                lru_width=None,
+            )
+            for i in range(config.num_hidden_layers)
+        ]
+        self.final_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def __call__(
+        self,
+        tokens,
+        mask: mx.array = None,
+        cache=None,
+    ):
+        x = self.embed_tokens(tokens)
+        if self.scale_by_sqrt_dim:
+            x = x * math.sqrt(x.shape[-1])
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        for i, block in enumerate(self.layers):
+            if block.temporal_block_type != "recurrent":
+                mask_cache = [cache[i]]
+
+        if mask is None:
+            mask = create_attention_mask(x, mask_cache)
+
+        for i, block in enumerate(self.layers):
+            x = block(x, mask=mask, cache=cache[i])
+
+        return self.final_norm(x)
+
+
+class Model(nn.Module):
+
+    def __init__(self, config):
+        self.args = config
+        self.model = Griffin(config)
+        self.model_type = config.model_type
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+    def __call__(self, tokens: mx.array, mask: mx.array = None, cache=None) -> mx.array:
+        """
+        Args:
+          tokens: Sequence of input tokens.
+        """
+        logits = self.model(tokens, mask=mask, cache=cache)
+        if "lm_head" in self:
+            logits = self.lm_head(logits)
+        else:
+            logits = self.model.embed_tokens.as_linear(logits)
+
+        c = self.args.logits_soft_cap
+        if c:
+            logits = mx.tanh(logits / c) * c
+        return logits
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def sanitize(self, weights):
+        for k, v in weights.items():
+            if "conv_1d.weight" in k and v.shape[-1] != 1:
+                weights[k] = v.moveaxis(2, 1)
+        if "lm_head.weight" not in weights:
+            self.pop("lm_head")
+        return weights
+
+    def make_cache(self):
+        cache = []
+        for layer in self.layers:
+            if layer.temporal_block_type == "recurrent":
+                cache.append(MambaCache())
+            else:
+                cache.append(RotatingKVCache(max_size=self.args.attention_window_size))
+        return cache

mlx_lm/models/rope_utils.py

@@ -0,0 +1,255 @@
+# Copyright © 2023-2024 Apple Inc.
+
+import math
+from typing import List, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class SuScaledRoPE(nn.Module):
+    def __init__(
+        self,
+        dims: int,
+        base: float = 10000.0,
+        max_position_embeddings: int = 131072,
+        original_max_position_embeddings: int = 4096,
+        short_factor: Union[List[float], float] = 1.0,
+        long_factor: Union[List[float], float] = 1.0,
+        short_mscale: float = None,
+        long_mscale: float = None,
+    ):
+        """
+        Su Scaled Rotary Embedding layer.
+
+        Args:
+            dims (int): The feature dimensions to be rotated.
+            base (int, optional): Base for the exponential scaling.
+            max_position_embeddings (int, optional): The maximum sequence
+              length that this model was trained with. This is used to determine
+              the size of the original RoPE embeddings when using long scaling.
+              Default: ``131072``.
+            original_max_position_embeddings (int, optional): The maximum
+              sequence length that this model was trained with. This is used to
+              determine the size of the original RoPE embeddings when using long
+              scaling. Default: ``4096``.
+            short_factor (float or list[float], optional): List of scaling
+              factors for sequences of length lesser than
+              ``original_max_position_embeddings``. Default: ``1.0``.
+            long_factor (float or list[float], optional): List of scaling
+              factors for sequences of length greater than
+              ``original_max_position_embeddings``.  Default: ``1.0``.
+            short_mscale (float, optional): Scale the input prior to embedding.
+            long_mscale (float, optional): Scale the input prior to embedding.
+        """
+        super().__init__()
+        freqs = base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
+        self._freqs = mx.array(long_factor, dtype=mx.float32) * freqs
+        self.original_max_position_embeddings = original_max_position_embeddings
+        self.scale = long_mscale or math.sqrt(
+            1
+            + math.log(max_position_embeddings / original_max_position_embeddings)
+            / math.log(original_max_position_embeddings)
+        )
+        self.dim = dims
+
+    def __call__(self, x, offset: int = 0):
+        x[..., : self.dim] = self.scale * x[..., : self.dim]
+        return mx.fast.rope(
+            x,
+            self.dim,
+            traditional=False,
+            base=None,
+            scale=1.0,
+            offset=offset,
+            freqs=self._freqs,
+        )
+
+
+class Llama3RoPE(nn.Module):
+
+    def __init__(
+        self,
+        dims: int,
+        max_position_embeddings: int = 2048,
+        traditional: bool = False,
+        base: float = 10000,
+        scaling_config: dict = None,
+    ):
+        super().__init__()
+        self.dims = dims
+        self.max_position_embeddings = max_position_embeddings
+        self.traditional = traditional
+
+        factor = scaling_config["factor"]
+        low_freq_factor = scaling_config.get("low_freq_factor", 1.0)
+        high_freq_factor = scaling_config.get("high_freq_factor", 4.0)
+        old_context_len = scaling_config.get(
+            "original_max_position_embeddings",
+            8192,
+        )
+
+        low_freq_wavelen = old_context_len / low_freq_factor
+        high_freq_wavelen = old_context_len / high_freq_factor
+
+        freqs = base ** (mx.arange(0, dims, 2) / dims)
+        wavelens = 2 * mx.pi * freqs
+
+        freqs = mx.where(wavelens > low_freq_wavelen, freqs * factor, freqs)
+        is_medium_freq = (wavelens > high_freq_wavelen) & (wavelens < low_freq_wavelen)
+        smooth_factors = (old_context_len / wavelens - low_freq_factor) / (
+            high_freq_factor - low_freq_factor
+        )
+        smooth_freqs = freqs / ((1 - smooth_factors) / factor + smooth_factors)
+        self._freqs = mx.where(is_medium_freq, smooth_freqs, freqs)
+
+    def extra_repr(self):
+        return (
+            f"{self.dims}, traditional={self.traditional}, "
+            f"max_position_embeddings={self.max_position_embeddings}"
+        )
+
+    def __call__(self, x, offset: int = 0):
+        return mx.fast.rope(
+            x,
+            self.dims,
+            traditional=self.traditional,
+            base=None,
+            scale=1.0,
+            offset=offset,
+            freqs=self._freqs,
+        )
+
+
+class YarnRoPE(nn.Module):
+    def __init__(
+        self,
+        dims,
+        traditional=False,
+        max_position_embeddings=2048,
+        base=10000,
+        scaling_factor=1.0,
+        original_max_position_embeddings=4096,
+        beta_fast=32,
+        beta_slow=1,
+        mscale=1,
+        mscale_all_dim=0,
+    ):
+        super().__init__()
+
+        def yarn_find_correction_dim(num_rotations):
+            return (
+                dims
+                * math.log(
+                    original_max_position_embeddings / (num_rotations * 2 * math.pi)
+                )
+            ) / (2 * math.log(base))
+
+        def yarn_find_correction_range():
+            low = math.floor(yarn_find_correction_dim(beta_fast))
+            high = math.ceil(yarn_find_correction_dim(beta_slow))
+            return max(low, 0), min(high, dims - 1)
+
+        def yarn_get_mscale(scale=1, mscale=1):
+            if scale <= 1:
+                return 1.0
+            return 0.1 * mscale * math.log(scale) + 1.0
+
+        def yarn_linear_ramp_mask(min_val, max_val, dim):
+            if min_val == max_val:
+                max_val += 0.001  # Prevent singularity
+
+            linear_func = (mx.arange(dim, dtype=mx.float32) - min_val) / (
+                max_val - min_val
+            )
+            return mx.clip(linear_func, 0, 1)
+
+        self.mscale = yarn_get_mscale(scaling_factor, mscale) / yarn_get_mscale(
+            scaling_factor, mscale_all_dim
+        )
+        freq_extra = base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
+        freq_inter = scaling_factor * base ** (
+            mx.arange(0, dims, 2, dtype=mx.float32) / dims
+        )
+        low, high = yarn_find_correction_range()
+        freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dims // 2)
+        self._freqs = (freq_inter * freq_extra) / (
+            freq_inter * freq_mask + freq_extra * (1 - freq_mask)
+        )
+        self.dims = dims
+        self.traditional = traditional
+
+    def __call__(self, x, offset=0):
+        if self.mscale != 1.0:
+            x[..., : self.dims] = self.mscale * x[..., : self.dims]
+        return mx.fast.rope(
+            x,
+            self.dims,
+            traditional=self.traditional,
+            base=None,
+            scale=1.0,
+            offset=offset,
+            freqs=self._freqs,
+        )
+
+
+def initialize_rope(
+    dims,
+    base,
+    traditional,
+    scaling_config: Optional[dict] = None,
+    max_position_embeddings: Optional[int] = None,
+):
+    if scaling_config is not None:
+        rope_type = scaling_config.get("type") or scaling_config.get(
+            "rope_type", "default"
+        )
+    else:
+        rope_type = "default"
+
+    if rope_type in ["default", "linear"]:
+        scale = 1 / scaling_config["factor"] if rope_type == "linear" else 1.0
+        return nn.RoPE(dims, traditional=traditional, base=base, scale=scale)
+
+    elif rope_type == "llama3":
+        return Llama3RoPE(
+            dims=dims,
+            max_position_embeddings=max_position_embeddings,
+            traditional=traditional,
+            base=base,
+            scaling_config=scaling_config,
+        )
+    elif rope_type == "yarn":
+        scaling_factor = scaling_config["factor"]
+        rope_kwargs = {
+            key: scaling_config[key]
+            for key in [
+                "original_max_position_embeddings",
+                "beta_fast",
+                "beta_slow",
+                "mscale",
+                "mscale_all_dim",
+            ]
+            if key in scaling_config
+        }
+        return YarnRoPE(
+            dims=dims,
+            max_position_embeddings=max_position_embeddings,
+            traditional=traditional,
+            base=base,
+            **rope_kwargs,
+        )
+    elif rope_type == "longrope":
+        return SuScaledRoPE(
+            dims=dims,
+            base=base,
+            max_position_embeddings=max_position_embeddings,
+            original_max_position_embeddings=scaling_config[
+                "original_max_position_embeddings"
+            ],
+            short_factor=scaling_config["short_factor"],
+            long_factor=scaling_config["long_factor"],
+        )
+
+    else:
+        raise ValueError(f"Unsupported RoPE type {rope_type}")

mlx_lm/models/smollm3.py

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

mlx_lm/models/stablelm.py

@@ -1,11 +1,12 @@
+# Copyright © 2023-2024 Apple Inc.
+
 import math
 from dataclasses import dataclass
-from typing import Tuple
 
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -119,8 +120,8 @@ class Attention(nn.Module):
 
         # Finally perform the attention computation
         scale = math.sqrt(1 / queries.shape[-1])
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=scale, mask=mask
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=scale, mask=mask
         ).astype(values.dtype)
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
         return self.o_proj(output)
@@ -196,12 +197,11 @@ class Model(nn.Module):
         self,
         x: mx.array,
         mask: mx.array = None,
-        cache: mx.array = None,
-    ) -> Tuple[mx.array, mx.array]:
-        mask = None
-        if x.shape[1] > 1:
-            mask = nn.MultiHeadAttention.create_additive_causal_mask(x.shape[1])
-            mask = mask.astype(x.dtype)
+        cache=None,
+    ) -> mx.array:
+
+        if mask is None:
+            mask = create_attention_mask(x, cache)
 
         y = self.model(x, mask, cache)
         return self.lm_head(y)
@@ -209,11 +209,3 @@ class Model(nn.Module):
     @property
     def layers(self):
         return self.model.layers
-
-    @property
-    def head_dim(self):
-        return self.args.hidden_size // self.args.num_attention_heads
-
-    @property
-    def n_kv_heads(self):
-        return self.args.num_key_value_heads

mlx_lm/models/starcoder2.py

@@ -1,10 +1,12 @@
+# Copyright © 2023-2024 Apple Inc.
+
 from dataclasses import dataclass
-from typing import Optional, Tuple
+from typing import Any, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
 @dataclass
@@ -43,7 +45,7 @@ class Attention(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         B, L, D = x.shape
 
@@ -62,8 +64,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
         )
 
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
@@ -98,7 +100,7 @@ class TransformerBlock(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Tuple[mx.array, mx.array]] = None,
+        cache: Optional[Any] = None,
     ) -> mx.array:
         r = self.self_attn(self.input_layernorm(x), mask, cache)
         h = x + r
@@ -123,14 +125,13 @@ class Starcoder2Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
         h = self.embed_tokens(inputs)
 
-        mask = None
-        if h.shape[1] > 1:
-            mask = nn.MultiHeadAttention.create_additive_causal_mask(h.shape[1])
-            mask = mask.astype(h.dtype)
+        if mask is None:
+            mask = create_attention_mask(h, cache)
 
         if cache is None:
             cache = [None] * len(self.layers)
@@ -153,9 +154,10 @@ class Model(nn.Module):
     def __call__(
         self,
         inputs: mx.array,
+        mask: mx.array = None,
         cache=None,
     ):
-        out = self.model(inputs, cache)
+        out = self.model(inputs, mask, cache)
         if self.args.tie_word_embeddings:
             out = self.model.embed_tokens.as_linear(out)
         else:
@@ -165,11 +167,3 @@ class Model(nn.Module):
     @property
     def layers(self):
         return self.model.layers
-
-    @property
-    def head_dim(self):
-        return self.args.hidden_size // self.args.num_attention_heads
-
-    @property
-    def n_kv_heads(self):
-        return self.args.num_key_value_heads

mlx_lm/models/su_rope.py

@@ -1,79 +0,0 @@
-import math
-from typing import List, Union
-
-import mlx.core as mx
-
-
-class SuScaledRotaryEmbedding:
-    def __init__(
-        self,
-        dims: int,
-        traditional: bool = False,
-        base: float = 10000.0,
-        scale: float = 1.0,
-        max_position_embeddings: int = 131072,
-        original_max_position_embeddings: int = 4096,
-        short_factor: Union[List[float], float] = 1.0,
-        long_factor: Union[List[float], float] = 1.0,
-    ):
-        """
-        Phi3Su Scaled Rotary Embedding layer for Phi-3 models.
-
-        Args:
-            dims (int): The feature dimensions to be rotated.
-            traditional (bool, optional): Unused. Default: ``False``.
-            base (int, optional): Base for the exponential scaling.
-            scale (float, optional): The scale used to scale the positions.
-              Default: ``1.0``.
-            max_position_embeddings (int, optional): The maximum sequence
-              length that this model was trained with. This is used to determine
-              the size of the original RoPE embeddings when using long scaling.
-              Default: ``131072``.
-            original_max_position_embeddings (int, optional): The maximum
-              sequence length that this model was trained with. This is used to
-              determine the size of the original RoPE embeddings when using long
-              scaling. Default: ``4096``.
-            short_factor (float or list[float], optional): List of scaling
-              factors for sequences of length lesser than
-              ``original_max_position_embeddings``. Default: ``1.0``.
-            long_factor (float or list[float], optional): List of scaling
-              factors for sequences of length greater than
-              ``original_max_position_embeddings``.  Default: ``1.0``.
-        """
-        self.inv_freq_short = 1.0 / (
-            mx.array(short_factor, dtype=mx.float32)
-            * base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
-        )
-        self.inv_freq_long = 1.0 / (
-            scale
-            * mx.array(long_factor, dtype=mx.float32)
-            * base ** (mx.arange(0, dims, 2, dtype=mx.float32) / dims)
-        )
-        self.original_max_position_embeddings = original_max_position_embeddings
-        self.scaling_factor = math.sqrt(
-            1
-            + math.log(max_position_embeddings / original_max_position_embeddings)
-            / math.log(original_max_position_embeddings)
-        )
-
-    def _get_cos_sin(self, offset, L):
-        position_ids = mx.arange(offset, offset + L, dtype=mx.float32)
-        inv_freq = (
-            self.inv_freq_long
-            if (offset + L) > self.original_max_position_embeddings
-            else self.inv_freq_short
-        )
-        freqs = position_ids[:, None] * inv_freq[None, :]
-        emb = mx.concatenate([freqs, freqs], axis=-1)
-        cos = mx.cos(emb) * self.scaling_factor
-        sin = mx.sin(emb) * self.scaling_factor
-        return cos, sin
-
-    def __call__(self, x, offset: int = 0):
-        def _rotate_half(_x):
-            midpoint = _x.shape[-1] // 2
-            x1, x2 = _x[..., :midpoint], _x[..., midpoint:]
-            return mx.concatenate([-x2, x1], axis=-1)
-
-        cos, sin = self._get_cos_sin(offset, x.shape[2])
-        return (x * cos) + (_rotate_half(x) * sin)