nits

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

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

.github/workflows/pull_request.yml

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

.github/workflows/release.yml

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

.gitignore

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

@@ -1,11 +0,0 @@
-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,25 +5,12 @@ with a short description of your contribution(s) below. For example:
 
 - Jane Smith: Added the `foo` example.
 
-MLX LM was developed with contributions from the following individuals:
+MLX Examples 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.
-- Gökdeniz Gülmez: Added support for the following architectures: 
-OpenBMB's `MiniCPM` and `MiniCPM3`, Kyutai's `Helium`, State-Space's `Mamba v1` and 
-`Mamba v2`, Z.ai & THUKEG's `GLM`, `GLM4`, Rednote `dots.llm1`, Baidu's `Ernie4.5 MoE`, 
-inclusionAI's `Bailing MoE e.g. Ling-family`, `Bailing MoE Linear e.g. Ling-Linear-family`, 
-Klear team - Kuaishou Technology's `Klear`, AI21 Lab's `Jamba` IBM's `Granite MoE`, 
-Meituan's `LongCat`, Nvidia's `Nemotron H`, Swiss-AI's `Apertus`, Nikity's `Lille130m`, 
-Alibaba Qwen's `Qwen3Next`, and Allenai's `OLMoE` and `Olmo 3`; 
-Helped add support for the following model architectures: 
-Alibaba Qwen's `Qwen3 & Qwen3MoE)`; Added support for the following training algorithms: 
-`Full Weight Fine-Tuning`, and the `Muon` optimizer; 
-Added support for the following other features: 
-`Multiple Optimizers to choose for training`, and `reporting training metrics to WandB (Weights & Biases)`.
-- Prince Canuma: Helped add support for the following model architectures:
-  HuggingFace's `Starcoder2`, Cohere's `Cohere (1 and 2)`, Alibaba Qwen's `Qwen (2, 3 and MoE)`, 
-  Microsoft's `Phi (3 and 3.5 MoE)`, `BitNet1.58`, Meta's `Llama (3 and 4)`, MinimaxAI's `MiniMax`, 
-  MoonshotAI's `Kimi-Linear`, LiquidAI's `LFM2` and `LFM2 MoE`, 
-  Google DeepMind's `Gemma 3`, TII's `Falcon H1` and InterLM's `InternLM 2.5`.
-- Ivan Fioravanti: Added support for the following architectures:
- ServiceNow-AI's `Apriel 1.5`, Tencent's `Hunyuan Dense V1` and `Hunyuan MoE V1`.
+- 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`.

CONTRIBUTING.md

@@ -1,54 +1,11 @@
 # 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.
 
-From this directory, do an editable install:
+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:
 
 ```shell
 pip install -e .
@@ -60,7 +17,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-lm/tree/main/mlx_lm/models)
+[`mlx_lm/models`](https://github.com/ml-explore/mlx-examples/tree/main/llms/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.
 
@@ -78,12 +35,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-lm/blob/main/mlx_lm/tuner/utils.py#L27-L60)
+[`mlx_lm/tuner/utils.py`](https://github.com/ml-explore/mlx-examples/blob/main/llms/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-lm/blob/main/tests/test_models.py).
+tests](https://github.com/ml-explore/mlx-examples/blob/main/llms/tests/test_models.py).
 
-You can run the tests with:
+From the `llms/` directory, you can run the tests with:
 
 ```shell
 python -m unittest discover tests/

MANIFEST.in

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

README.md

@@ -1,17 +1,4 @@
-## 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`
+## Generate Text with LLMs and MLX
 
 The easiest way to get started is to install the `mlx-lm` package:
 
@@ -27,36 +14,11 @@ pip install mlx-lm
 conda install -c conda-forge mlx-lm
 ```
 
-### Quick Start
+The `mlx-lm` package also has:
 
-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
-```
-
-The default model for generation and chat is
-`mlx-community/Llama-3.2-3B-Instruct-4bit`.  You can specify any MLX-compatible
-model with the `--model` flag. Thousands are available in the
-[MLX Community](https://huggingface.co/mlx-community) Hugging Face
-organization.
+- [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)
 
 ### Python API
 
@@ -67,14 +29,7 @@ from mlx_lm import load, generate
 
 model, tokenizer = load("mlx-community/Mistral-7B-Instruct-v0.3-4bit")
 
-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)
+response = generate(model, tokenizer, prompt="hello", verbose=True)
 ```
 
 To see a description of all the arguments you can do:
@@ -83,16 +38,10 @@ 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. Check out the [batch generation
-example](https://github.com/ml-explore/mlx-lm/tree/main/mlx_lm/examples/batch_generate_response.py)
-to see how to efficiently generate continuations for a batch of prompts.
-
 The `mlx-lm` package also comes with functionality to quantize and optionally
 upload models to the Hugging Face Hub.
 
-You can convert models using the Python API:
+You can convert models in the Python API with:
 
 ```python
 from mlx_lm import convert
@@ -115,10 +64,8 @@ To see a description of all the arguments you can do:
 
 #### Streaming
 
-For streaming generation, use the `stream_generate` function. This yields
-a generation response object.
-
-For example,
+For streaming generation, use the `stream_generate` function. This returns a
+generator object which streams the output text. For example,
 
 ```python
 from mlx_lm import load, stream_generate
@@ -128,28 +75,11 @@ model, tokenizer = load(repo)
 
 prompt = "Write a story about Einstein"
 
-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)
+for t in stream_generate(model, tokenizer, prompt, max_tokens=512):
+    print(t, 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:
@@ -190,65 +120,46 @@ 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
 
-`mlx-lm` supports thousands of LLMs available on the Hugging Face Hub. If the
-model you want to run is not supported, file an
-[issue](https://github.com/ml-explore/mlx-lm/issues/new) or better yet, submit
-a pull request. Many supported models are available in various quantization
-formats in the [MLX Community](https://huggingface.co/mlx-community) Hugging
-Face organization.
+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,
+submit a pull request.
 
-For some models the tokenizer may require you to enable the `trust_remote_code`
-option. You can do this by passing `--trust-remote-code` in the command line.
-If you don't specify the flag explicitly, you will be prompted to trust remote
-code in the terminal when running the model. 
+Here are a few examples of Hugging Face models that work with this example:
 
-Tokenizer options can also be set in the Python API. For example:
+- [mistralai/Mistral-7B-v0.1](https://huggingface.co/mistralai/Mistral-7B-v0.1)
+- [meta-llama/Llama-2-7b-hf](https://huggingface.co/meta-llama/Llama-2-7b-hf)
+- [deepseek-ai/deepseek-coder-6.7b-instruct](https://huggingface.co/deepseek-ai/deepseek-coder-6.7b-instruct)
+- [01-ai/Yi-6B-Chat](https://huggingface.co/01-ai/Yi-6B-Chat)
+- [microsoft/phi-2](https://huggingface.co/microsoft/phi-2)
+- [mistralai/Mixtral-8x7B-Instruct-v0.1](https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1)
+- [Qwen/Qwen-7B](https://huggingface.co/Qwen/Qwen-7B)
+- [pfnet/plamo-13b](https://huggingface.co/pfnet/plamo-13b)
+- [pfnet/plamo-13b-instruct](https://huggingface.co/pfnet/plamo-13b-instruct)
+- [stabilityai/stablelm-2-zephyr-1_6b](https://huggingface.co/stabilityai/stablelm-2-zephyr-1_6b)
+- [internlm/internlm2-7b](https://huggingface.co/internlm/internlm2-7b)
+
+Most
+[Mistral](https://huggingface.co/models?library=transformers,safetensors&other=mistral&sort=trending),
+[Llama](https://huggingface.co/models?library=transformers,safetensors&other=llama&sort=trending),
+[Phi-2](https://huggingface.co/models?library=transformers,safetensors&other=phi&sort=trending),
+and
+[Mixtral](https://huggingface.co/models?library=transformers,safetensors&other=mixtral&sort=trending)
+style models should work out of the box.
+
+For some models (such as `Qwen` and `plamo`) the tokenizer requires you to
+enable the `trust_remote_code` option. You can do this by passing
+`--trust-remote-code` in the command line. If you don't specify the flag
+explicitly, you will be prompted to trust remote code in the terminal when
+running the model. 
+
+For `Qwen` models you must also specify the `eos_token`. You can do this by
+passing `--eos-token "<|endoftext|>"` in the command
+line. 
+
+These options can also be set in the Python API. For example:
 
 ```python
 model, tokenizer = load(
@@ -256,28 +167,3 @@ 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/BENCHMARKS.md

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

mlx_lm/LEARNED_QUANTS.md

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

mlx_lm/LORA.md

@@ -26,12 +26,6 @@ LoRA (QLoRA).[^qlora] LoRA fine-tuning works with the following model families:
 
 ## Run
 
-First, make sure you have the training dependenices installed:
-
-```shell
-pip install "mlx-lm[train]"
-```
-
 The main command is `mlx_lm.lora`. To see a full list of command-line options run:
 
 ```shell
@@ -63,9 +57,6 @@ 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).
@@ -76,31 +67,12 @@ 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 learned weights are saved in `adapters/`.
-You can specify the output location with `--adapter-path`.
+By default, the adapter config and 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 using `--report-to wandb`, or
-to SwanLab using `--report-to swanlab`. Make sure to install the required
-packages beforehand: `pip install wandb` or `pip install swanlab`. You can
-enable both tracking tools simultaneously by separating them with a comma, for
-example: `--report-to wandb,swanlab`.
-
-To specify a project name for the logging tracker, use `--project-name <YOUR
-PROJECT NAME>`.
-
-#### Prompt Masking
-
-The default training computes a loss for every token in the sample. You can
-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:
@@ -146,7 +118,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 `fused_model/`. All of these are configurable.
+model in the path `lora_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
@@ -169,7 +141,7 @@ mlx_lm.fuse \
     --export-gguf
 ```
 
-This will save the GGUF model in `fused_model/ggml-model-f16.gguf`. You
+This will save the GGUF model in `lora_fused_model/ggml-model-f16.gguf`. You
 can specify the file name with `--gguf-path`.
 
 ## Data
@@ -179,173 +151,59 @@ 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 `chat`, `tools`, `completions`, and `text`
-data formats. Here are examples of these formats:
+Currently, `*.jsonl` files support three data formats: `chat`,
+`completions`, and `text`. Here are three 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": "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"]
-                }
-            }
-        }
-    ]
+  "messages": [
+    {
+      "role": "system",
+      "content": "You are a helpful assistant."
+    },
+    {
+      "role": "user",
+      "content": "Hello."
+    },
+    {
+      "role": "assistant",
+      "content": "How can I assistant you today."
+    }
+  ]
 }
 ```
 
-
-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.
 
-> [!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:
+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:
 
 ```text
 <|im_start|>system
@@ -371,12 +229,9 @@ of memory. Here are some tips to reduce memory use should you need to do so:
 
 2. Try using a smaller batch size with `--batch-size`. The default is `4` so
    setting this to `2` or `1` will reduce memory consumption. This may slow
-   things down a little, but will also reduce the memory use. You can increase
-   the effective batch size without increasing the memory use by accumulating
-   gradients using `--grad-accumulation-steps <N>` which will accumulate the
-   gradient of `<N>` batches before updating the parameters.
+   things down a little, but will also reduce the memory use.
 
-3. Reduce the number of layers to fine-tune with `--num-layers`. The default
+3. Reduce the number of layers to fine-tune with `--lora-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.
@@ -398,8 +253,8 @@ mlx_lm.lora \
     --model mistralai/Mistral-7B-v0.1 \
     --train \
     --batch-size 1 \
-    --num-layers 4 \
-    --data  mlx-community/wikisql
+    --lora-layers 4 \
+    --data wikisql
 ```
 
 The above command on an M1 Max with 32 GB runs at about 250
@@ -408,5 +263,4 @@ 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

@@ -0,0 +1,50 @@
+# 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 mlx-community/Mistral-7B-Instruct-v0.3-4bit
+mlx_lm.server --model mistralai/Mistral-7B-Instruct-v0.1
 ```
 
 This will start a text generation server on port `8080` of the `localhost`
@@ -50,28 +50,22 @@ 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. These are
+- `stop`: (Optional) An array of strings or a single string. Thesse 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 `512`.
+  to generate. Defaults to `100`.
 
 - `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 `0.0`.
+  Defaults to `1.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`.
 
@@ -79,65 +73,4 @@ 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`.
-
-- `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.
+  values. Defaults to `None`.

mlx_lm/UPLOAD.md

@@ -0,0 +1,37 @@
+### 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,20 +1,4 @@
 # Copyright © 2023-2024 Apple Inc.
 
-import os
-
-from ._version import __version__
-
-os.environ["TRANSFORMERS_NO_ADVISORY_WARNINGS"] = "1"
-
-from .convert import convert
-from .generate import batch_generate, generate, stream_generate
-from .utils import load
-
-__all__ = [
-    "__version__",
-    "convert",
-    "batch_generate",
-    "generate",
-    "stream_generate",
-    "load",
-]
+from .utils import convert, generate, load, stream_generate
+from .version import __version__

mlx_lm/__main__.py

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

mlx_lm/_version.py

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

mlx_lm/benchmark.py

@@ -1,126 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-import argparse
-
-import mlx.core as mx
-
-from mlx_lm import batch_generate, load, stream_generate
-from mlx_lm.generate import DEFAULT_MODEL
-from mlx_lm.utils import pipeline_load
-
-
-def setup_arg_parser():
-    """Set up and return the argument parser."""
-    parser = argparse.ArgumentParser(description="LLM benchmarking script")
-    parser.add_argument(
-        "--model",
-        type=str,
-        help=(
-            "The path to the local model directory or Hugging Face repo. "
-            f"If no model is specified, then {DEFAULT_MODEL} is used."
-        ),
-        default=None,
-    )
-    parser.add_argument(
-        "--prompt-tokens",
-        "-p",
-        default=512,
-        help="Length of prompt",
-        type=int,
-    )
-    parser.add_argument(
-        "--generation-tokens",
-        "-g",
-        default=1024,
-        help="Length of completion",
-        type=int,
-    )
-    parser.add_argument(
-        "--batch-size",
-        "-b",
-        default=1,
-        help="Batch size",
-        type=int,
-    )
-    parser.add_argument(
-        "--num-trials",
-        "-n",
-        default=5,
-        help="Number of timing trials",
-        type=int,
-    )
-    return parser
-
-
-def main():
-    parser = setup_arg_parser()
-    args = parser.parse_args()
-    mx.random.seed(0)
-
-    group = mx.distributed.init()
-    rank = group.rank()
-
-    def rprint(*args, **kwargs):
-        if rank == 0:
-            print(*args, **kwargs)
-
-    model_path = args.model or DEFAULT_MODEL
-
-    if group.size() > 1:
-        model, tokenizer, config = pipeline_load(args.model, return_config=True)
-    else:
-        model, tokenizer, config = load(
-            args.model, return_config=True, tokenizer_config={"trust_remote_code": True}
-        )
-
-    # Empty to avoid early stopping
-    tokenizer._eos_token_ids = {}
-
-    prompt_tokens = args.prompt_tokens
-    generation_tokens = args.generation_tokens
-    batch_size = args.batch_size
-    vocab_size = config.get("vocab_size") or config["text_config"]["vocab_size"]
-    prompts = mx.random.randint(0, vocab_size, (batch_size, prompt_tokens)).tolist()
-    prompt = prompts[0]
-
-    def single_bench():
-        for response in stream_generate(
-            model, tokenizer, prompt, max_tokens=generation_tokens
-        ):
-            pass
-        return response
-
-    def batch_bench():
-        return batch_generate(
-            model, tokenizer, prompts, max_tokens=generation_tokens
-        ).stats
-
-    if batch_size == 1:
-        _bench = single_bench
-    else:
-        _bench = batch_bench
-
-    rprint("Running warmup..")
-    _bench()
-
-    report_keys = ["prompt_tps", "generation_tps", "peak_memory"]
-    rprint(f"Timing with {prompt_tokens=}, {generation_tokens=}, {batch_size=}.")
-    responses = []
-    for i in range(args.num_trials):
-        response = _bench()
-        responses.append(response)
-        results = [(k, getattr(response, k)) for k in report_keys]
-        results = [f"{k}={v:.3f}" for k, v in results]
-        rprint(f"Trial {i+1}:  " + ", ".join(results))
-
-    def avg(k):
-        vals = (getattr(response, k) for response in responses)
-        return sum(vals) / args.num_trials
-
-    results = [(k, avg(k)) for k in report_keys]
-    results = [f"{k}={v:.3f}" for k, v in results]
-    rprint(f"Averages: " + ", ".join(results))
-
-
-if __name__ == "__main__":
-    main()

mlx_lm/cache_prompt.py

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

@@ -1,149 +0,0 @@
-# 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(
-        "--trust-remote-code",
-        action="store_true",
-        help="Enable trusting remote code for tokenizer",
-    )
-    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",
-    )
-    parser.add_argument(
-        "--system-prompt",
-        default=None,
-        help="System prompt to be used for the chat template",
-    )
-    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 if args.trust_remote_code else None
-        },
-    )
-
-    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 = []
-        if args.system_prompt is not None:
-            messages.append({"role": "system", "content": args.system_prompt})
-        messages.append({"role": "user", "content": query})
-        prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
-        for response in stream_generate(
-            model,
-            tokenizer,
-            prompt,
-            max_tokens=args.max_tokens,
-            sampler=make_sampler(
-                args.temp,
-                args.top_p,
-                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,171 +1,8 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import argparse
-from pathlib import Path
-from typing import Callable, Optional, Union
 
-import mlx.core as mx
-import mlx.nn as nn
-from mlx.utils import tree_map_with_path
-
-from .utils import (
-    dequantize_model,
-    load,
-    quantize_model,
-    save,
-    upload_to_hub,
-)
-
-
-def mixed_quant_predicate_builder(
-    recipe: str, model: nn.Module, group_size: int = 64
-) -> Callable[[str, nn.Module, dict], Union[bool, dict]]:
-    high_bits = 6
-
-    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(f"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,
-    ) -> 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 or "v_a_proj" in path or "v_b_proj" in path
-        ) and use_more_bits:
-            return {"group_size": group_size, "bits": high_bits}
-        if "down_proj" in path and use_more_bits:
-            return {"group_size": group_size, "bits": high_bits}
-        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,
-    q_mode: str = "affine",
-    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, tokenizer, config = load(
-        hf_path,
-        revision=revision,
-        return_config=True,
-        tokenizer_config={"trust_remote_code": trust_remote_code},
-        lazy=True,
-    )
-
-    if isinstance(quant_predicate, str):
-        quant_predicate = mixed_quant_predicate_builder(
-            quant_predicate, model, q_group_size
-        )
-
-    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,
-            mode=q_mode,
-            quant_predicate=quant_predicate,
-        )
-
-    if dequantize:
-        print("[INFO] Dequantizing")
-        config.pop("quantization", None)
-        config.pop("quantization_config", None)
-        model = dequantize_model(model)
-
-    save(
-        mlx_path,
-        hf_path,
-        model,
-        tokenizer,
-        config,
-    )
-
-    if upload_repo is not None:
-        upload_to_hub(mlx_path, upload_repo)
+from .utils import convert
 
 
 def configure_parser() -> argparse.ArgumentParser:
@@ -192,26 +29,12 @@ def configure_parser() -> argparse.ArgumentParser:
     parser.add_argument(
         "--q-bits", help="Bits per weight for quantization.", type=int, default=4
     )
-    parser.add_argument(
-        "--q-mode",
-        help="The quantization mode.",
-        type=str,
-        default="affine",
-        choices=["affine", "mxfp4"],
-    )
-    parser.add_argument(
-        "--quant-predicate",
-        help=f"Mixed-bit quantization recipe.",
-        choices=QUANT_RECIPES,
-        type=str,
-        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.",
+        help="Type to save the parameters, ignored if -q is given.",
         type=str,
-        choices=MODEL_CONVERSION_DTYPES,
-        default=None,
+        choices=["float16", "bfloat16", "float32"],
+        default="float16",
     )
     parser.add_argument(
         "--upload-repo",
@@ -226,12 +49,6 @@ def configure_parser() -> argparse.ArgumentParser:
         action="store_true",
         default=False,
     )
-    parser.add_argument(
-        "--trust-remote-code",
-        help="Trust remote code when loading tokenizer.",
-        action="store_true",
-        default=False,
-    )
     return parser
 
 
@@ -242,8 +59,4 @@ 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

@@ -1,509 +0,0 @@
-# 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, Callable, 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 lm_eval.models import huggingface
-from tqdm import tqdm
-
-from .generate import batch_generate
-from .models.cache import make_prompt_cache
-from .sample_utils import make_sampler
-from .utils import load
-
-DEFAULT_MAX_TOKENS = 8192
-
-
-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 _lstrip(s, pattern):
-    """Truncate the prefix of the string after the first occurrence of pattern."""
-    if (idx := s.find(pattern)) != -1:
-        return s[idx + len(pattern) :]
-    return s
-
-
-def _pad_inputs(inputs):
-    lengths = np.array([len(x) for x in inputs])
-    maxlen = lengths.max()
-    padded = np.stack(
-        [np.pad(x, (0, maxlen - len(x))) for x in inputs],
-        axis=0,
-    )
-    return mx.array(padded), mx.array(lengths)
-
-
-def chat_template_fn(**extra_kwargs):
-    def apply_chat_template(self, chat_history, add_generation_prompt=True) -> str:
-        return self.tokenizer.apply_chat_template(
-            chat_history,
-            tokenize=False,
-            add_generation_prompt=add_generation_prompt,
-            continue_final_message=not add_generation_prompt,
-            **extra_kwargs,
-        )
-
-    return apply_chat_template
-
-
-@register_model("mlxlm")
-class MLXLM(LM):
-
-    tokenizer_name = huggingface.HFLM.tokenizer_name
-    apply_chat_template = chat_template_fn()
-
-    def __init__(
-        self,
-        path_or_hf_repo: str,
-        max_tokens: Optional[int] = None,
-        use_chat_template: Optional[bool] = None,
-        trust_remote_code: bool = False,
-        sampler: Optional[Callable[[mx.array], mx.array]] = None,
-    ) -> None:
-        super().__init__()
-        tokenizer_config = {"trust_remote_code": True if trust_remote_code else None}
-        self._model, self.tokenizer = load(
-            path_or_hf_repo, tokenizer_config=tokenizer_config
-        )
-        self._max_tokens = max_tokens
-        self._batch_size = 8
-        self.use_chat_template = use_chat_template
-        if use_chat_template is None:
-            self.use_chat_template = self.tokenizer.chat_template is not None
-        self._sampler = sampler
-
-    def _process_prompt(self, prompt, step_size: int = 2048):
-        prompt = mx.array(prompt)[None]
-        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)
-        offset = 0
-        scores, is_greedy = [], []
-        for i in range(0, inputs.shape[1], step_size):
-            inp = inputs[:, i : i + step_size]
-            T = inp.shape[1]
-
-            logits = self._model(inp, cache=cache)
-            log_probs = nn.log_softmax(logits.astype(mx.float32))
-
-            score = mx.take_along_axis(
-                log_probs, targets[:, i : i + step_size, mx.newaxis], axis=-1
-            )[..., 0]
-
-            ig = targets[:, i : i + step_size] == mx.argmax(logits, axis=-1)
-            ig = mx.where(mx.arange(offset, T + offset) < lengths[:, None], ig, False)
-
-            mx.eval(score, ig)
-            mx.clear_cache()
-
-            is_greedy.append(ig)
-            scores.append(score)
-            offset += T
-
-        scores = mx.concatenate(scores, axis=1)
-        is_greedy = mx.concatenate(is_greedy, axis=1)
-
-        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.append(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
-            max_tokens = self._max_tokens or DEFAULT_MAX_TOKENS
-            truncation = max(0, max_completed_l - max_tokens - 1)
-            orig_prefix_l = len(prefix)
-            prefix_l = max(len(prefix) - truncation, 0)
-            prefix = prefix[len(prefix) - prefix_l :]
-
-            # 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()
-
-        if long_completions > 0:
-            logging.info(
-                f"Prefix eliminated for {long_completions} requests with "
-                + "completion longer than context."
-            )
-
-        # All gather the results across nodes
-        num_results = len(requests)
-        per_group = mx.distributed.all_max(len(scores), stream=mx.cpu).item()
-        scores = scores + [0] * (per_group - len(scores))
-        is_greedy = is_greedy + [False] * (per_group - len(is_greedy))
-        scores = mx.array(scores)
-        is_greedy = mx.array(is_greedy)
-        scores = mx.distributed.all_gather(scores, stream=mx.cpu)
-        is_greedy = mx.distributed.all_gather(is_greedy, stream=mx.cpu)
-        mx.eval(scores, is_greedy)
-
-        # Arrange the indices to match the scores from each node and then
-        # inverse sort the scores
-        all_indices = []
-        for rank in range(group.size()):
-            rank_indices = [
-                idx for question in indices[rank :: group.size()] for idx in question
-            ]
-            rank_indices += [num_results] * (per_group - len(rank_indices))
-            all_indices.extend(rank_indices)
-        inv_sort = mx.argsort(mx.array(all_indices))
-        scores = scores[:num_results][inv_sort]
-        is_greedy = is_greedy[:num_results][inv_sort]
-
-        return list(zip(scores.tolist(), is_greedy.tolist()))
-
-    def loglikelihood_rolling(self, requests) -> list[float]:
-        """Compute full log-likelihood of a string, with no truncation, for perplexity computation
-        - 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(inputs), self._batch_size)):
-            batch = inputs[i : i + self._batch_size]
-            scores, lengths, _ = self._score_fn(batch)
-            mask = mx.arange(scores.shape[-1]) < lengths[:, None]
-            all_scores.extend((mask * scores).sum(axis=-1).tolist())
-
-        return all_scores
-
-    def generate_until(self, requests) -> list[str]:
-        """Generate greedily until a stopping sequence
-        :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.
-        """
-        group = mx.distributed.init()
-
-        # split data accross ranks
-        total_requests = len(requests)
-        requests = requests[group.rank() :: group.size()]
-
-        logging.info("Generating continuation for %d sequences." % len(requests))
-        contexts, options = zip(*[req.args for req in requests])
-        # The second element of the tuple contains:
-        # {'do_sample': False, 'until': ['\n\n'], 'temperature': 0}
-
-        # Tokenize all contexts
-        contexts = [
-            self.tokenizer.encode(
-                context, add_special_tokens=not self.use_chat_template
-            )
-            for context in contexts
-        ]
-
-        # TODO consider multi-token, per-prompt stop conditions
-        max_tokens = [
-            self._max_tokens or opt.get("max_gen_tokens", DEFAULT_MAX_TOKENS)
-            for opt in options
-        ]
-
-        completions = batch_generate(
-            model=self._model,
-            tokenizer=self.tokenizer,
-            prompts=contexts,
-            max_tokens=max_tokens,
-            verbose=True,
-            sampler=self._sampler,
-        ).texts
-
-        for e, (text, opt) in enumerate(zip(completions, options)):
-            completions[e] = _rstrip_until(text, opt["until"])
-            if self.tokenizer.has_thinking:
-                completions[e] = _lstrip(text, self.tokenizer.think_end)
-
-        # Gather the completions
-        if group.size() > 1:
-            with mx.stream(mx.cpu):
-                pad_to = (total_requests + group.size() - 1) // group.size()
-                pad = pad_to - len(completions)
-                completions = [list(c.encode("utf-8")) for c in completions]
-                max_len = mx.array(max(len(c) for c in completions))
-                max_len = mx.distributed.all_max(max_len).item()
-                lengths = mx.array([len(c) for c in completions] + [0] * pad)
-                completions = mx.array(
-                    [c + [0] * (max_len - len(c)) for c in completions]
-                    + [[0] * max_len] * pad,
-                    mx.uint8,
-                )
-                completions = (
-                    mx.distributed.all_gather(completions[None])
-                    .swapaxes(0, 1)
-                    .flatten(0, 1)
-                    .tolist()
-                )
-                lengths = (
-                    mx.distributed.all_gather(lengths[None])
-                    .swapaxes(0, 1)
-                    .flatten(0, 1)
-                    .tolist()
-                )
-                completions = completions[:total_requests]
-                lengths = lengths[:total_requests]
-                completions = [
-                    bytearray(c[:l]).decode() for c, l in zip(completions, lengths)
-                ]
-
-        return completions
-
-
-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 number of tokens to generate. When set, this value takes"
-        " precedence over task specific defaults.",
-        default=None,
-    )
-    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="{}",
-    )
-    parser.add_argument(
-        "--confirm-run-unsafe-code",
-        action="store_true",
-        help="Confirm that you want to run tasks that execute untrusted code.",
-        default=False,
-    )
-    parser.add_argument(
-        "--trust-remote-code",
-        action="store_true",
-        help="Enable trusting remote code for tokenizer",
-    )
-    parser.add_argument("--temp", type=float, default=0.0, help="Sampling temperature")
-    parser.add_argument("--top-p", type=float, default=1.0, help="Sampling top-p")
-    parser.add_argument("--top-k", type=int, default=0, help="Sampling top-k")
-    args = parser.parse_args()
-
-    output_dir = Path(args.output_dir)
-    output_dir.mkdir(parents=True, exist_ok=True)
-
-    # Silence tokenizer warnings
-    os.environ["TOKENIZERS_PARALLELISM"] = "false"
-
-    mx.random.seed(args.seed)
-
-    # Initialize the communication if in distributed mode
-    world = mx.distributed.init()
-    mx.eval(mx.distributed.all_sum(1, stream=mx.cpu))
-    if world.size() > 1 and world.rank() == 0:
-        print(f"Evaluating with {world.size()} nodes")
-
-    sampler = make_sampler(
-        temp=args.temp,
-        top_p=args.top_p,
-        top_k=args.top_k,
-    )
-    lm = MLXLM(
-        args.model,
-        max_tokens=args.max_tokens,
-        use_chat_template=args.apply_chat_template,
-        trust_remote_code=args.trust_remote_code,
-        sampler=sampler,
-    )
-    MLXLM.apply_chat_template = chat_template_fn(**args.chat_template_args)
-
-    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,
-        confirm_run_unsafe_code=args.confirm_run_unsafe_code,
-    )
-
-    file_keys = ["eval", args.model.replace("/", "_"), version("lm_eval")]
-    if args.num_shots is not None:
-        file_keys += [f"{args.num_shots:02d}"]
-    file_keys += args.tasks
-    filename = "_".join(file_keys)
-    if world.rank() == 0:
-        output_path = output_dir / filename
-        output_path.write_text(json.dumps(results["results"], indent=4))
-        print("Results:")
-        for result in results["results"].values():
-            print(json.dumps(result, indent=4))

mlx_lm/examples/batch_generate_response.py

@@ -1,51 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-from mlx_lm import batch_generate, load
-
-# Specify the checkpoint
-checkpoint = "mlx-community/Llama-3.2-3B-Instruct-4bit"
-
-# Load the corresponding model and tokenizer
-model, tokenizer = load(path_or_hf_repo=checkpoint)
-
-# A batch of prompts
-prompts = [
-    "Write a story about Einstein.",
-    "Why is the sky blue?",
-    "What time is it?",
-    "How tall is Mt Everest?",
-]
-
-# Apply the chat template and encode to tokens
-prompts = [
-    tokenizer.apply_chat_template(
-        [{"role": "user", "content": p}],
-        add_generation_prompt=True,
-    )
-    for p in prompts
-]
-
-# Set `verbose=True` to see generation statistics
-result = batch_generate(
-    model, tokenizer, prompts, verbose=False, return_prompt_caches=True
-)
-print(result.texts[-1])
-
-prompts = [
-    "Could you summarize that?",
-    "And what about the sea?",
-    "Try again?",
-    "And Mt Olympus?",
-]
-prompts = [
-    tokenizer.apply_chat_template(
-        [{"role": "user", "content": p}],
-        add_generation_prompt=True,
-    )
-    for p in prompts
-]
-
-result = batch_generate(
-    model, tokenizer, prompts, verbose=False, prompt_caches=result.caches
-)
-print(result.texts[-1])

mlx_lm/examples/chat.py

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

@@ -1,33 +0,0 @@
-# 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,29 +1,16 @@
 # The path to the local model directory or Hugging Face repo.
-model: "mlx-community/Llama-3.2-1B-Instruct-bf16"
-
+model: "mlx_model"
 # 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: "mlx-community/WikiSQL"
+data: "/path/to/training/data"
 
 # The PRNG seed
 seed: 0
 
 # Number of layers to fine-tune
-num_layers: 16
+lora_layers: 16
 
 # Minibatch size.
 batch_size: 4
@@ -37,19 +24,12 @@ val_batches: 25
 # Adam learning rate.
 learning_rate: 1e-5
 
-# Services to report logs to (comma-separated): wandb, swanlab, or both ('wandb,swanlab').
-# report_to: wandb,swanlab
-# project_name: "Your-awesome-mlx-project-name"
-
 # Number of training steps between loss reporting.
 steps_per_report: 10
 
 # Number of training steps between validations.
 steps_per_eval: 200
 
-# Number of micro-steps to accumulate before each optimizer update.
-grad_accumulation_steps: 1
-
 # Load path to resume training with the given adapter weights.
 resume_adapter_file: null
 
@@ -71,6 +51,9 @@ 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.
@@ -86,10 +69,3 @@ 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

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

@@ -1,75 +0,0 @@
-# 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 mlx.core as mx
-
-from mlx_lm import stream_generate
-from mlx_lm.utils import pipeline_load
-
-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 = pipeline_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

@@ -1,73 +0,0 @@
-# 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,13 +1,19 @@
 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 .utils import (
-    dequantize_model,
-    load,
-    save,
+    fetch_from_hub,
+    get_model_path,
+    save_config,
+    save_weights,
     upload_to_hub,
 )
 
@@ -23,7 +29,7 @@ def parse_arguments() -> argparse.Namespace:
     )
     parser.add_argument(
         "--save-path",
-        default="fused_model",
+        default="lora_fused_model",
         help="The path to save the fused model.",
     )
     parser.add_argument(
@@ -32,6 +38,12 @@ 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.",
@@ -39,8 +51,8 @@ def parse_arguments() -> argparse.Namespace:
         default=None,
     )
     parser.add_argument(
-        "--dequantize",
-        help="Generate a dequantized model.",
+        "--de-quantize",
+        help="Generate a de-quantized model.",
         action="store_true",
     )
     parser.add_argument(
@@ -61,33 +73,40 @@ def main() -> None:
     print("Loading pretrained model")
     args = parse_arguments()
 
-    model, tokenizer, config = load(
-        args.model, adapter_path=args.adapter_path, return_config=True
-    )
+    model_path = get_model_path(args.model)
+    model, config, tokenizer = fetch_from_hub(model_path)
+
+    model.freeze()
+    model = apply_lora_layers(model, args.adapter_path)
 
     fused_linears = [
-        (n, m.fuse(dequantize=args.dequantize))
+        (n, m.to_linear())
         for n, m in model.named_modules()
-        if hasattr(m, "fuse")
+        if isinstance(m, (LoRASwitchLinear, LoRALinear, DoRALinear))
     ]
 
-    if fused_linears:
-        model.update_modules(tree_unflatten(fused_linears))
+    model.update_modules(tree_unflatten(fused_linears))
 
-    if args.dequantize:
-        print("Dequantizing model")
+    if args.de_quantize:
+        print("De-quantizing model")
         model = dequantize(model)
-        config.pop("quantization", None)
+
+    weights = dict(tree_flatten(model.parameters()))
 
     save_path = Path(args.save_path)
-    save(
-        save_path,
-        args.model,
-        model,
-        tokenizer,
-        config,
-        donate_model=False,
-    )
+
+    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")
 
     if args.export_gguf:
         model_type = config["model_type"]
@@ -95,16 +114,18 @@ def main() -> None:
             raise ValueError(
                 f"Model type {model_type} not supported for GGUF conversion."
             )
-        weights = dict(tree_flatten(model.parameters()))
-        convert_to_gguf(save_path, weights, config, str(save_path / args.gguf_path))
+        convert_to_gguf(model_path, weights, config, str(save_path / args.gguf_path))
 
     if args.upload_repo is not None:
-        upload_to_hub(args.save_path, args.upload_repo)
+        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)
 
 
 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/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]
+            token_text = reverse_vocab[token_id].encode("utf-8")
             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(r"<0x[0-9A-Fa-f]{2}>", token_text):
+        if re.fullmatch(rb"<0x[0-9A-Fa-f]{2}>", token_text):
             return TokenType.BYTE
         return TokenType.CONTROL if token_id in special_ids else TokenType.NORMAL
 
@@ -77,12 +77,14 @@ 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], "", self.special_ids)
+                toktype = self.get_token_type(
+                    self.specials[text], b"", self.special_ids
+                )
                 score = self.get_token_score(self.specials[text])
             else:
                 toktype = TokenType.USER_DEFINED
                 score = -1000.0
-            yield text, score, toktype
+            yield text.encode("utf-8"), score, toktype
 
     def has_newline_token(self):
         return "<0x0A>" in self.tokenizer.vocab or "\n" in self.tokenizer.vocab
@@ -241,18 +243,15 @@ 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)
-    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["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
+    )
 
     metadata = {k: v for k, v in metadata.items() if v is not None}
     return metadata

mlx_lm/lora.py

@@ -1,24 +1,23 @@
+# Copyright © 2024 Apple Inc.
+
 import argparse
 import math
-import os
 import re
 import types
-import warnings
 from pathlib import Path
 
-import mlx.core as mx
 import mlx.nn as nn
 import mlx.optimizers as optim
 import numpy as np
 import yaml
 
-from .tuner.callbacks import get_reporting_callbacks
-from .tuner.datasets import CacheDataset, load_dataset
+from .tokenizer_utils import TokenizerWrapper
+from .tuner.datasets import 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
@@ -40,20 +39,11 @@ yaml_loader.add_implicit_resolver(
 )
 
 CONFIG_DEFAULTS = {
-    "model": "Qwen/Qwen3-0.6b",
+    "model": "mlx_model",
     "train": False,
-    "fine_tune_type": "lora",
-    "optimizer": "adam",
-    "optimizer_config": {
-        "adam": {},
-        "adamw": {},
-        "muon": {},
-        "sgd": {},
-        "adafactor": {},
-    },
-    "data": "mlx-community/WikiSQL",
+    "data": "data/",
     "seed": 0,
-    "num_layers": 16,
+    "lora_layers": 16,
     "batch_size": 4,
     "iters": 1000,
     "val_batches": 25,
@@ -66,14 +56,9 @@ CONFIG_DEFAULTS = {
     "test": False,
     "test_batches": 500,
     "max_seq_length": 2048,
-    "config": None,
-    "grad_checkpoint": False,
-    "grad_accumulation_steps": 1,
     "lr_schedule": None,
-    "lora_parameters": {"rank": 8, "dropout": 0.0, "scale": 20.0},
-    "mask_prompt": False,
-    "report_to": None,
-    "project_name": None,
+    "lora_parameters": {"rank": 8, "alpha": 16, "dropout": 0.0, "scale": 10.0},
+    "use_dora": False,
 }
 
 
@@ -81,7 +66,6 @@ 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.",
     )
 
@@ -95,32 +79,10 @@ def build_parser():
     parser.add_argument(
         "--data",
         type=str,
-        help=(
-            "Directory with {train, valid, test}.jsonl files or the name "
-            "of a Hugging Face dataset (e.g., 'mlx-community/wikisql')"
-        ),
+        help="Directory with {train, valid, test}.jsonl files",
     )
     parser.add_argument(
-        "--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", "muon", "sgd", "adafactor"],
-        default=None,
-        help="Optimizer to use for training: adam, adamw, sgd, or adafactor.",
-    )
-    parser.add_argument(
-        "--mask-prompt",
-        action="store_true",
-        help="Mask the prompt in the loss when training",
-        default=None,
-    )
-    parser.add_argument(
-        "--num-layers",
+        "--lora-layers",
         type=int,
         help="Number of layers to fine-tune. Default is 16, use -1 for all.",
     )
@@ -142,20 +104,15 @@ def build_parser():
         type=int,
         help="Number of training steps between validations.",
     )
-    parser.add_argument(
-        "--grad-accumulation-steps",
-        type=int,
-        help="Number of steps to accumulate before each optimizer update.",
-    )
     parser.add_argument(
         "--resume-adapter-file",
         type=str,
-        help="Load path to resume training from the given fine-tuned weights.",
+        help="Load path to resume training with the given adapters.",
     )
     parser.add_argument(
         "--adapter-path",
         type=str,
-        help="Save/load path for the fine-tuned weights.",
+        help="Save/load path for the adapters.",
     )
     parser.add_argument(
         "--save-every",
@@ -181,7 +138,7 @@ def build_parser():
     parser.add_argument(
         "-c",
         "--config",
-        type=str,
+        default=None,
         help="A YAML configuration file with the training options",
     )
     parser.add_argument(
@@ -190,63 +147,36 @@ 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(
-        "--report-to",
-        type=str,
-        default=None,
-        help="Services to report logs to ('wandb', 'swanlab', or 'wandb,swanlab').",
+        "--use-dora", action="store_true", default=None, help="Use DoRA to finetune."
     )
-    parser.add_argument(
-        "--project-name",
-        type=str,
-        default=None,
-        help="Project name for logging. Defaults to the name of the root directory.",
-    )
-    parser.add_argument("--seed", type=int, help="The PRNG seed")
     return parser
 
 
 def train_model(
     args,
     model: nn.Module,
+    tokenizer: TokenizerWrapper,
     train_set,
     valid_set,
     training_callback: TrainingCallback = None,
 ):
-    mx.random.seed(args.seed)
+    # Freeze all layers
     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."
-        )
 
-    if args.fine_tune_type == "full":
-        for l in model.layers[-max(args.num_layers, 0) :]:
-            l.unfreeze()
+    # Convert linear layers to lora layers and unfreeze in the process
+    linear_to_lora_layers(model, args.lora_layers, args.lora_parameters)
 
-        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
+    # Resume training the given adapters.
     if args.resume_adapter_file is not None:
-        print(f"Loading fine-tuned weights from {args.resume_adapter_file}")
+        print(f"Loading pretrained adapters 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")
 
@@ -261,44 +191,33 @@ def train_model(
         adapter_file=adapter_file,
         max_seq_length=args.max_seq_length,
         grad_checkpoint=args.grad_checkpoint,
-        grad_accumulation_steps=args.grad_accumulation_steps,
     )
 
-    # Initialize the selected optimizer
-    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
-    elif optimizer_name == "muon":
-        opt_class = optim.Muon
-    elif optimizer_name == "sgd":
-        opt_class = optim.SGD
-    elif optimizer_name == "adafactor":
-        opt_class = optim.Adafactor
-    else:
-        raise ValueError(f"Unsupported optimizer: {optimizer_name}")
-
-    opt = opt_class(learning_rate=lr, **optimizer_config)
-
+    model.train()
+    opt = optim.Adam(
+        learning_rate=(
+            build_schedule(args.lr_schedule) if args.lr_schedule else args.learning_rate
+        )
+    )
     # Train model
     train(
         model=model,
+        tokenizer=tokenizer,
         args=training_args,
         optimizer=opt,
-        train_dataset=CacheDataset(train_set),
-        val_dataset=CacheDataset(valid_set),
+        train_dataset=train_set,
+        val_dataset=valid_set,
         training_callback=training_callback,
     )
 
 
-def evaluate_model(args, model: nn.Module, test_set):
+def evaluate_model(args, model: nn.Module, tokenizer: TokenizerWrapper, test_set):
+    model.eval()
+
     test_loss = evaluate(
         model=model,
-        dataset=CacheDataset(test_set),
+        dataset=test_set,
+        tokenizer=tokenizer,
         batch_size=args.batch_size,
         num_batches=args.test_batches,
         max_seq_length=args.max_seq_length,
@@ -311,15 +230,9 @@ def evaluate_model(args, model: nn.Module, test_set):
 
 def run(args, training_callback: TrainingCallback = None):
     np.random.seed(args.seed)
-    training_callback = get_reporting_callbacks(
-        args.report_to,
-        project_name=args.project_name,
-        log_dir=args.adapter_path,
-        config=vars(args),
-    )
 
     print("Loading pretrained model")
-    model, tokenizer = load(args.model, tokenizer_config={"trust_remote_code": True})
+    model, tokenizer = load(args.model)
 
     print("Loading datasets")
     train_set, valid_set, test_set = load_dataset(args, tokenizer)
@@ -327,21 +240,20 @@ 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 != "":
-            load_adapters(model, args.adapter_path)
+            apply_lora_layers(model, args.adapter_path)
 
     elif args.train:
         print("Training")
-        train_model(args, model, train_set, valid_set, training_callback)
+        train_model(args, model, tokenizer, 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, test_set)
+        evaluate_model(args, model, tokenizer, test_set)
 
 
 def main():
-    os.environ["TOKENIZERS_PARALLELISM"] = "true"
     parser = build_parser()
     args = parser.parse_args()
     config = args.config
@@ -363,8 +275,4 @@ 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,37 +2,23 @@ import argparse
 from typing import List, Union
 
 from huggingface_hub import scan_cache_dir
-
-
-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)
+from transformers.commands.user import tabulate
 
 
 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", "")
-    full_message = f"{message} (y/n) "
+    n = ("n", "no", "0")
+    all_values = y + n + ("",)
+    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: yes/no/y/n or empty for no")
+        print(f"Invalid input. Must be one of {all_values}")
 
 
 def main():
@@ -57,7 +43,9 @@ def main():
     args = parser.parse_args()
 
     if args.scan:
-        print(f'Scanning Hugging Face cache for models with pattern "{args.pattern}".')
+        print(
+            "Scanning Hugging Face cache for models with" f'pattern "{args.pattern}".'
+        )
         hf_cache_info = scan_cache_dir()
         print(
             tabulate(
@@ -98,46 +86,36 @@ 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(
-                "\nAre you sure you want to delete these models?"
-            )
+            confirmed = ask_for_confirmation(f"Confirm deletion ?")
             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("\nModel(s) deleted successfully.")
+                print("Model(s) deleted.")
             else:
-                print("\nDeletion cancelled - no changes made.")
+                print("Deletion is cancelled. Do nothing.")
         else:
-            print(f'No models found matching pattern "{args.pattern}"')
+            print(f"No models found.")
 
 
 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

@@ -0,0 +1,172 @@
+# 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/Klear.py

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

mlx_lm/models/afm7.py

@@ -1,394 +0,0 @@
-# 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, dequantize: 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 dequantize:
-            fused_linear = fused_linear.to_quantized(linear.group_size, linear.bits)
-
-        return fused_linear
-
-    def __call__(self, x):
-        dt = x.dtype
-        y = self.linear(x)
-        x = self.dropout(x)
-        z = [(x @ a) @ b for a, b in zip(self.lora_a, self.lora_b)]
-        return tuple(yi + (self.scale * zi).astype(dt) for yi, zi in zip(y, z))
-
-
-class FusedQuantizedLinear(nn.QuantizedLinear):
-    def __init__(self, input_dims, output_dims, group_size: int = 64, bits: int = 4):
-        *indices, output_dims = accumulate(output_dims)
-        self.indices = indices
-        super().__init__(
-            input_dims, output_dims, bias=False, group_size=group_size, bits=bits
-        )
-
-    @property
-    def input_dims(self):
-        return self.scales.shape[-1] * self.group_size
-
-    @property
-    def output_dims(self):
-        indices = [0] + self.indices + [self.weight.shape[0]]
-        return [indices[i] - indices[i - 1] for i in range(1, len(indices))]
-
-    def __call__(self, x):
-        x = super().__call__(x)
-        return x.split(self.indices, axis=-1)
-
-    def to_lora(self, r: int = 8, dropout: float = 0.0, scale: float = 20.0):
-        lora_lin = FusedLoRALinear(self.input_dims, self.output_dims, r, dropout, scale)
-        lora_lin.linear = self
-        return lora_lin
-
-
-class FusedLinear(nn.Linear):
-    def __init__(self, input_dims, output_dims):
-        *indices, output_dims = accumulate(output_dims)
-        self.indices = indices
-        super().__init__(input_dims, output_dims, bias=False)
-
-    @property
-    def input_dims(self):
-        return self.weight.shape[-1]
-
-    @property
-    def output_dims(self):
-        indices = [0] + self.indices + [self.weight.shape[0]]
-        return [indices[i] - indices[i - 1] for i in range(1, len(indices))]
-
-    def __call__(self, x):
-        x = super().__call__(x)
-        return x.split(self.indices, axis=-1)
-
-    def to_quantized(self, group_size: int = 64, bits: int = 4):
-        input_dims = self.input_dims
-        output_dims = self.output_dims
-        ql = FusedQuantizedLinear(input_dims, output_dims, group_size, bits)
-        ql.weight, ql.scales, ql.biases = mx.quantize(self.weight, group_size, bits)
-
-        return ql
-
-    def to_lora(self, r: int = 8, dropout: float = 0.0, scale: float = 20.0):
-        lora_lin = FusedLoRALinear(self.input_dims, self.output_dims, r, dropout, scale)
-        lora_lin.linear = self
-        return lora_lin
-
-
-@partial(mx.compile, shapeless=True)
-def fake_8bit_quant(x, scale):
-    dt = x.dtype
-    x = x.astype(mx.float32)
-    x = (x / scale).round()
-    x = mx.clip(x, -128, 127)
-    return (x * scale).astype(dt)
-
-
-class Attention(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-
-        dim = args.hidden_dim
-        self.n_heads = n_heads = args.num_heads
-        self.n_kv_heads = n_kv_heads = args.num_kv_heads
-        self.head_dim = head_dim = args.hidden_dim // n_heads
-        self.scale = head_dim**-0.5
-
-        qkv_dim = (n_heads + 2 * n_kv_heads) * head_dim
-        self.qkv_proj = FusedLinear(
-            dim, [n_heads * head_dim] + 2 * [n_kv_heads * head_dim]
-        )
-        self.out_proj = nn.Linear(dim, dim, bias=False)
-        self.rope = initialize_rope(
-            self.head_dim,
-            args.rope_theta,
-            True,
-        )
-        self.q_norm = nn.RMSNorm(head_dim)
-        self.k_norm = nn.RMSNorm(head_dim)
-        self.quant_key_scale = mx.array(1.0)
-        self.quant_value_scale = mx.array(1.0)
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-    ) -> mx.array:
-        B, L, D = x.shape
-
-        # Get the queries, keys and values
-        queries, keys, values = self.qkv_proj(x)
-
-        # Prepare the queries, keys and values for the attention computation
-        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
-        keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
-        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
-
-        if cache is not None:
-            queries = self.q_norm(self.rope(queries, offset=cache.offset))
-            keys = self.k_norm(self.rope(keys, offset=cache.offset))
-            keys = fake_8bit_quant(keys, self.quant_key_scale)
-            values = fake_8bit_quant(values, self.quant_value_scale)
-            keys, values = cache.update_and_fetch(keys, values)
-        else:
-            queries = self.q_norm(self.rope(queries))
-            keys = self.k_norm(self.rope(keys))
-            keys = fake_8bit_quant(keys, self.quant_key_scale)
-            values = fake_8bit_quant(values, self.quant_value_scale)
-
-        output = scaled_dot_product_attention(
-            queries, keys, values, cache=cache, scale=self.scale, mask=mask
-        )
-
-        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
-        return self.out_proj(output)
-
-
-class KVReuseAttention(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-
-        dim = args.hidden_dim
-        self.n_heads = n_heads = args.num_heads
-        self.head_dim = head_dim = args.hidden_dim // n_heads
-        self.scale = head_dim**-0.5
-
-        self.q_proj = nn.Linear(dim, dim, bias=False)
-        self.out_proj = nn.Linear(dim, dim, bias=False)
-        self.rope = initialize_rope(
-            self.head_dim,
-            args.rope_theta,
-            True,
-        )
-        self.q_norm = nn.RMSNorm(head_dim)
-
-    def __call__(
-        self,
-        x: mx.array,
-        keys: mx.array,
-        values: mx.array,
-        mask: Optional[mx.array] = None,
-    ) -> mx.array:
-        B, L, D = x.shape
-        _, _, S, _ = keys.shape
-
-        queries = self.q_proj(x)
-        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
-        queries = self.q_norm(self.rope(queries, offset=S - L))
-
-        output = scaled_dot_product_attention(
-            queries, keys, values, cache=None, scale=self.scale, mask=mask
-        )
-
-        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
-        return self.out_proj(output)
-
-
-@partial(mx.compile, shapeless=True)
-def _swiglu(g, x):
-    return nn.silu(g) * x
-
-
-class MLP(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-
-        dim = args.hidden_dim
-        hidden_dim = int(dim * args.hidden_dim_scale_factor)
-
-        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
-        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
-        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
-
-    def __call__(self, x) -> mx.array:
-        g = self.gate_proj(x)
-        x = self.up_proj(x)
-        return self.down_proj(_swiglu(g, x))
-
-
-class TransformerBlock(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.self_attn = Attention(args)
-        self.mlp = MLP(args)
-        self.input_layernorm = nn.RMSNorm(args.hidden_dim, eps=args.rms_norm_eps)
-        self.post_attention_layernorm = nn.RMSNorm(
-            args.hidden_dim, eps=args.rms_norm_eps
-        )
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-    ) -> mx.array:
-        r = self.self_attn(self.input_layernorm(x), mask, cache)
-        h = x + r
-        r = self.mlp(self.post_attention_layernorm(h))
-        out = h + r
-        return out
-
-
-class KVReuseTransformerBlock(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.self_attn = KVReuseAttention(args)
-        self.mlp = MLP(args)
-        self.input_layernorm = nn.RMSNorm(args.hidden_dim, eps=args.rms_norm_eps)
-        self.post_attention_layernorm = nn.RMSNorm(
-            args.hidden_dim, eps=args.rms_norm_eps
-        )
-
-    def __call__(
-        self,
-        x: mx.array,
-        keys: mx.array,
-        values: mx.array,
-        mask: Optional[mx.array] = None,
-    ) -> mx.array:
-        r = self.self_attn(self.input_layernorm(x), keys, values, mask)
-        h = x + r
-        r = self.mlp(self.post_attention_layernorm(h))
-        out = h + r
-        return out
-
-
-class AFMModel(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.vocab_size = args.vocab_size
-
-        self.embedding = nn.Embedding(args.vocab_size, args.hidden_dim)
-        self.layers = [
-            TransformerBlock(args)
-            for _ in range(args.num_layers - args.num_kv_reuse_layers)
-        ]
-        self.kv_reuse_layers = [
-            KVReuseTransformerBlock(args) for _ in range(args.num_kv_reuse_layers)
-        ]
-        self.output_norm = nn.RMSNorm(args.hidden_dim, eps=args.rms_norm_eps)
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        cache=None,
-    ):
-        h = self.embedding(inputs)
-
-        if cache is None:
-            cache = [None] * len(self.layers)
-            cache[-1] = ConcatenateKVCache()
-
-        mask = create_attention_mask(h, cache[0])
-
-        for layer, c in zip(self.layers, cache):
-            h = layer(h, mask, cache=c)
-
-        keys, values = cache[-1].state
-        for layer in self.kv_reuse_layers:
-            h = layer(h, keys, values, mask)
-
-        return self.output_norm(h)
-
-
-class Model(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.model_type = args.model_type
-        self.model = AFMModel(args)
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        cache=None,
-    ):
-        out = self.model(inputs, cache)
-        out = self.model.embedding.as_linear(out)
-        return out
-
-    def make_cache(self):
-        return [KVCache() for _ in range(len(self.model.layers))]
-
-    @property
-    def layers(self):
-        return self.model.layers + self.model.kv_reuse_layers

mlx_lm/models/afmoe.py

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

mlx_lm/models/apertus.py

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

mlx_lm/models/baichuan_m1.py

@@ -1,250 +0,0 @@
-# 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 None:
-            cache = (None, None)
-
-        if cache[0] is not None:
-            offset = cache[1].offset
-            last_k, last_v = cache[0][0], cache[0][1]
-        else:
-            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[0] 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)
-
-        self.sliding_window = config.sliding_window
-        self.first_swa_idx = None
-        if config.sliding_window_layers:
-            self.first_swa_idx = config.sliding_window_layers[0]
-
-        self.first_global_idx = None
-        self.swa_layers = set(config.sliding_window_layers)
-        for i in range(config.num_hidden_layers):
-            if i in self.swa_layers:
-                continue
-            self.first_global_idx = i
-            break
-
-    def __call__(self, inputs: mx.array, cache: Any = None) -> mx.array:
-        x = self.embed_tokens(inputs)
-
-        if cache is None:
-            cache = [(None, None)] * len(self.layers)
-
-        if self.first_global_idx is None:
-            c_global = None
-        else:
-            c_global = cache[self.first_global_idx][1]
-
-        if self.first_swa_idx is None:
-            c_swa = None
-        else:
-            c_swa = cache[self.first_swa_idx][1]
-
-        global_mask = create_attention_mask(x, c_global)
-        swa_mask = create_attention_mask(x, c_swa, window_size=self.sliding_window)
-
-        for l, (layer, c) in enumerate(zip(self.layers, cache)):
-            mask = swa_mask if l in self.swa_layers else global_mask
-            x = layer(x, mask, c)
-        return self.norm(x)
-
-
-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, cache: Any = None) -> mx.array:
-        outputs = self.model(inputs, cache)
-        return self.lm_head(outputs)
-
-    @property
-    def layers(self) -> List[nn.Module]:
-        return self.model.layers

mlx_lm/models/bailing_moe.py

@@ -1,405 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-from dataclasses import dataclass
-from functools import partial
-from typing import Any, Dict, Optional, Union
-
-import mlx.core as mx
-import mlx.nn as nn
-
-from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
-from .rope_utils import initialize_rope
-from .switch_layers import SwitchGLU
-
-
-@dataclass
-class ModelArgs(BaseModelArgs):
-    model_type: str
-    hidden_size: int
-    intermediate_size: int
-    max_position_embeddings: int
-    moe_intermediate_size: int
-    num_experts: int
-    num_shared_experts: int
-    norm_topk_prob: bool
-    num_attention_heads: int
-    num_experts_per_tok: int
-    num_hidden_layers: int
-    num_key_value_heads: int
-    rms_norm_eps: float
-    rope_theta: float
-    vocab_size: int
-    first_k_dense_replace: int
-    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
-    use_bias: bool = False
-    use_qkv_bias: bool = False
-    norm_head: bool = False
-    norm_softmax: bool = False
-    use_qk_norm: bool = False
-    tie_word_embeddings: bool = False
-    partial_rotary_factor: float = 1.0
-    rotary_dim: Optional[int] = None
-    moe_router_enable_expert_bias: bool = False
-    moe_router_enable_routed_scaling: bool = True
-    routed_scaling_factor: float = 1.0
-    score_function: str = "softmax"
-    n_group: int = 1
-    topk_group: int = 4
-    moe_shared_expert_intermediate_size: Optional[int] = None
-    moe_router_enable_shared_expert: bool = True
-
-
-@partial(mx.compile, shapeless=True)
-def swiglu(gate, up):
-    return nn.silu(gate) * up
-
-
-@partial(mx.compile, shapeless=True)
-def aggregate_expert_outputs(expert_outputs, scores):
-    return (
-        (expert_outputs * scores[..., None]).sum(axis=-2).astype(expert_outputs.dtype)
-    )
-
-
-class BailingMoeMLP(nn.Module):
-    def __init__(self, args: ModelArgs, intermediate_size: Optional[int] = None):
-        super().__init__()
-        self.intermediate_size = (
-            intermediate_size
-            if intermediate_size is not None
-            else args.intermediate_size
-        )
-
-        self.gate_proj = nn.Linear(
-            args.hidden_size, self.intermediate_size, bias=args.use_bias
-        )
-        self.down_proj = nn.Linear(
-            self.intermediate_size, args.hidden_size, bias=args.use_bias
-        )
-        self.up_proj = nn.Linear(
-            args.hidden_size, self.intermediate_size, bias=args.use_bias
-        )
-
-    def __call__(self, x) -> mx.array:
-        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
-
-
-class BailingMoeAttention(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.use_qk_norm = args.use_qk_norm
-        self.num_attention_heads = args.num_attention_heads
-        self.num_key_value_heads = args.num_key_value_heads
-        self.head_dim = args.hidden_size // self.num_attention_heads
-        self.scale = self.head_dim**-0.5
-
-        self.query_key_value = nn.Linear(
-            args.hidden_size,
-            (self.num_attention_heads + 2 * self.num_key_value_heads) * self.head_dim,
-            bias=args.use_qkv_bias,
-        )
-        self.dense = nn.Linear(
-            self.num_attention_heads * self.head_dim,
-            args.hidden_size,
-            bias=args.use_bias,
-        )
-
-        if args.use_qk_norm:
-            self.key_layernorm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
-            self.query_layernorm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
-
-        if (rope_dim := args.rotary_dim) is None:
-            rope_dim = int(self.head_dim * args.partial_rotary_factor)
-        self.rope = initialize_rope(
-            rope_dim,
-            args.rope_theta,
-            traditional=False,
-            scaling_config=args.rope_scaling,
-            max_position_embeddings=args.max_position_embeddings,
-        )
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-    ) -> mx.array:
-        B, L, D = x.shape
-
-        qkv = self.query_key_value(x)
-
-        q_size = self.num_attention_heads * self.head_dim
-        kv_size = self.num_key_value_heads * self.head_dim
-        q, k, v = mx.split(qkv, [q_size, q_size + kv_size], axis=-1)
-
-        queries = q.reshape(B, L, self.num_attention_heads, self.head_dim).transpose(
-            0, 2, 1, 3
-        )
-        keys = k.reshape(B, L, self.num_key_value_heads, self.head_dim).transpose(
-            0, 2, 1, 3
-        )
-        values = v.reshape(B, L, self.num_key_value_heads, self.head_dim).transpose(
-            0, 2, 1, 3
-        )
-
-        if self.use_qk_norm:
-            queries = self.query_layernorm(queries)
-            keys = self.key_layernorm(keys)
-
-        if cache is not None:
-            queries = self.rope(queries, offset=cache.offset)
-            keys = self.rope(keys, offset=cache.offset)
-            keys, values = cache.update_and_fetch(keys, values)
-        else:
-            queries = self.rope(queries)
-            keys = self.rope(keys)
-
-        output = scaled_dot_product_attention(
-            queries, keys, values, cache=cache, scale=self.scale, mask=mask
-        )
-
-        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
-        return self.dense(output)
-
-
-@mx.compile
-def group_expert_select(
-    gates,
-    e_score_correction_bias,
-    top_k,
-    n_group,
-    topk_group,
-    routed_scaling_factor,
-    norm_topk_prob,
-    score_function,
-):
-
-    in_type = gates.dtype
-    if score_function == "sigmoid":
-        scores = mx.sigmoid(gates.astype(mx.float32))
-    else:
-        scores = mx.softmax(gates.astype(mx.float32), axis=-1)
-    orig_scores = scores
-    if e_score_correction_bias is not None:
-        scores = scores + e_score_correction_bias
-    if n_group > 1:
-        scores = mx.unflatten(scores, axis=-1, shape=(n_group, -1))
-        group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True)
-        k = n_group - topk_group
-        group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :]
-        scores = mx.put_along_axis(
-            scores, mx.stop_gradient(group_idx), mx.array(0.0, scores.dtype), axis=-2
-        )
-        scores = mx.flatten(scores, -2, -1)
-
-    k = top_k
-    inds = mx.argpartition(scores, kth=-k, axis=-1)[..., -k:]
-    scores = mx.take_along_axis(orig_scores, inds, axis=-1)
-    if top_k > 1 and norm_topk_prob:
-        denominator = scores.sum(axis=-1, keepdims=True) + 1e-20
-        scores = scores / denominator
-    scores = scores * routed_scaling_factor
-
-    return inds, scores.astype(in_type)
-
-
-class BailingMoeGate(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.norm_topk_prob = args.norm_topk_prob
-
-        self.top_k = args.num_experts_per_tok
-        self.n_group = args.n_group
-        self.topk_group = args.topk_group
-        self.routed_scaling_factor = args.routed_scaling_factor
-        self.enable_routed_scaling = args.moe_router_enable_routed_scaling
-
-        self.gate_proj = nn.Linear(args.hidden_size, args.num_experts, bias=False)
-        self.expert_bias = (
-            mx.zeros((args.num_experts,))
-            if args.moe_router_enable_expert_bias
-            else None
-        )
-        self.score_function = args.score_function
-
-    def __call__(self, x):
-        return group_expert_select(
-            self.gate_proj(x),
-            self.expert_bias,
-            self.top_k,
-            self.n_group,
-            self.topk_group,
-            self.routed_scaling_factor,
-            self.norm_topk_prob,
-            self.score_function,
-        )
-
-
-class BailingMoeSparseMoeBlock(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.num_experts_per_tok = args.num_experts_per_tok
-        self.switch_mlp = SwitchGLU(
-            args.hidden_size,
-            args.moe_intermediate_size,
-            args.num_experts,
-            bias=args.use_bias,
-        )
-        self.gate = BailingMoeGate(args)
-        shared_dim = (
-            args.moe_shared_expert_intermediate_size or args.moe_intermediate_size
-        )
-        self.shared_experts = (
-            BailingMoeMLP(
-                args=args,
-                intermediate_size=shared_dim * args.num_shared_experts,
-            )
-            if args.num_shared_experts > 0 and args.moe_router_enable_shared_expert
-            else None
-        )
-
-    def __call__(self, x):
-        topk_idx, topk_weight = self.gate(x)
-        out = self.switch_mlp(x, topk_idx)
-        out = aggregate_expert_outputs(out, topk_weight)
-        if self.shared_experts is not None:
-            out = out + self.shared_experts(x)
-        return out
-
-
-class BailingMoeDecoderLayer(nn.Module):
-    def __init__(self, args: ModelArgs, layer_idx: int):
-        super().__init__()
-        self.attention = BailingMoeAttention(args)
-
-        self.mlp = (
-            BailingMoeSparseMoeBlock(args)
-            if (
-                args.num_experts is not None and layer_idx >= args.first_k_dense_replace
-            )
-            else BailingMoeMLP(args)
-        )
-        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
-        self.post_attention_layernorm = nn.RMSNorm(
-            args.hidden_size, eps=args.rms_norm_eps
-        )
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-    ) -> mx.array:
-        r = self.attention(self.input_layernorm(x), mask, cache)
-        h = x + r
-        r = self.mlp(self.post_attention_layernorm(h))
-        return h + r
-
-
-class BailingMoeModel(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.word_embeddings = nn.Embedding(args.vocab_size, args.hidden_size)
-        self.layers = [
-            BailingMoeDecoderLayer(args, layer_idx=i)
-            for i in range(args.num_hidden_layers)
-        ]
-        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        cache: Optional[Any] = None,
-    ):
-        h = self.word_embeddings(inputs)
-
-        if cache is None:
-            cache = [None] * len(self.layers)
-
-        mask = create_attention_mask(h, cache[0])
-
-        for layer, c in zip(self.layers, cache):
-            h = layer(h, mask, c)
-
-        return self.norm(h)
-
-
-class Model(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.norm_head = args.norm_head
-        self.model_type = args.model_type
-        self.model = BailingMoeModel(args)
-        if not args.tie_word_embeddings:
-            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        cache=None,
-    ):
-        out = self.model(inputs, cache)
-        if self.args.tie_word_embeddings:
-            out = self.model.word_embeddings.as_linear(out)
-        else:
-            out = self.lm_head(out)
-        return out
-
-    def sanitize(self, weights):
-        if self.args.tie_word_embeddings:
-            weights.pop("lm_head.weight", None)
-
-        if self.norm_head:
-            w = weights["lm_head.weight"]
-            dtype = w.dtype
-            weight_norm = (
-                mx.linalg.norm(w.astype(mx.float32), axis=0, keepdims=True) + 1e-7
-            )
-            weights["lm_head.weight"] = (w / weight_norm).astype(dtype)
-
-        for l in range(self.args.num_hidden_layers):
-            prefix = f"model.layers.{l}"
-
-            if l >= self.args.first_k_dense_replace:
-                for m in ["gate_proj", "down_proj", "up_proj"]:
-                    for k in ["weight", "scales", "biases"]:
-                        if f"{prefix}.mlp.experts.0.{m}.{k}" in weights:
-                            to_join = [
-                                weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}")
-                                for e in range(self.args.num_experts)
-                            ]
-                            weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(
-                                to_join
-                            )
-
-                if f"{prefix}.mlp.gate.weight" in weights:
-                    gate_weight = weights.pop(f"{prefix}.mlp.gate.weight")
-                    weights[f"{prefix}.mlp.gate.gate_proj.weight"] = gate_weight
-
-                if f"{prefix}.mlp.gate.bias" in weights:
-                    gate_bias = weights.pop(f"{prefix}.mlp.gate.bias")
-                    weights[f"{prefix}.mlp.gate.gate_proj.bias"] = gate_bias
-
-        return weights
-
-    @property
-    def quant_predicate(self):
-        def predicate(path, _):
-            if path.endswith("mlp.gate.gate_proj"):
-                return {"group_size": 64, "bits": 8}
-            return True
-
-        return predicate
-
-    @property
-    def cast_predicate(self):
-        def predicate(k):
-            return "expert_bias" not in k
-
-        return predicate
-
-    @property
-    def layers(self):
-        return self.model.layers

mlx_lm/models/bailing_moe_linear.py

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

mlx_lm/models/base.py

@@ -1,11 +1,46 @@
-# 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, :]
 
 
 @dataclass
@@ -19,119 +54,3 @@ class BaseModelArgs:
                 if k in inspect.signature(cls).parameters
             }
         )
-
-
-def create_causal_mask(
-    N: int,
-    offset: int = 0,
-    window_size: Optional[int] = None,
-    right_padding: Optional[mx.array] = None,
-    left_padding: Optional[mx.array] = None,
-):
-    rinds = mx.arange(offset + N)
-    linds = mx.arange(offset, offset + N) if offset else rinds
-    linds = linds[:, None]
-    rinds = rinds[None]
-    mask = linds >= rinds
-    if window_size is not None:
-        mask = mask & (linds < rinds + window_size)
-    if right_padding is not None:
-        mask = mask & (rinds < mx.expand_dims((offset + N) - right_padding, (1, 2, 3)))
-    if left_padding is not None:
-        mask = mask & (mx.expand_dims(left_padding, (1, 2, 3)) <= rinds)
-    return mask
-
-
-def create_attention_mask(
-    h, cache=None, window_size: Optional[int] = None, return_array: bool = False
-):
-    N = h.shape[1]
-    if cache and hasattr(cache, "make_mask"):
-        return cache.make_mask(N, return_array=return_array, window_size=window_size)
-    if N == 1:
-        return None
-    if return_array or (window_size and N > window_size):
-        return create_causal_mask(N, window_size=window_size)
-    return "causal"
-
-
-def create_ssm_mask(h, cache=None):
-    if cache and hasattr(cache, "make_mask"):
-        return cache.make_mask(h.shape[1])
-    return None
-
-
-def quantized_scaled_dot_product_attention(
-    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],
-    sinks: Optional[mx.array] = None,
-) -> mx.array:
-    if hasattr(cache, "bits"):
-        if sinks is not None:
-            raise ValueError("Quantized SDPA does not support attention sinks.")
-        return quantized_scaled_dot_product_attention(
-            queries,
-            keys,
-            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,
-            sinks=sinks,
-        )

mlx_lm/models/bitlinear_layers.py

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

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

mlx_lm/models/cohere.py

@@ -1,12 +1,10 @@
-# Copyright © 2023-2024 Apple Inc.
-
 from dataclasses import dataclass
-from typing import Any, Optional
+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 .base import BaseModelArgs
 
 
 @dataclass
@@ -69,7 +67,7 @@ class Attention(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = None,
     ) -> mx.array:
         B, L, D = x.shape
 
@@ -93,8 +91,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = scaled_dot_product_attention(
-            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        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)
@@ -129,7 +127,7 @@ class TransformerBlock(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = None,
     ) -> mx.array:
         h = self.input_layernorm(x)
         attn_h = self.self_attn(h, mask, cache)
@@ -159,11 +157,14 @@ class CohereModel(nn.Module):
     ):
         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 cache is None:
             cache = [None] * len(self.layers)
 
-        mask = create_attention_mask(h, cache[0])
-
         for layer, c in zip(self.layers, cache):
             h = layer(h, mask, c)
 
@@ -190,3 +191,11 @@ 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

@@ -1,216 +0,0 @@
-# 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)
-
-        # 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.window_size = args.sliding_window
-        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,
-        cache=None,
-    ):
-        h = self.embed_tokens(inputs)
-
-        if cache is None:
-            cache = [None] * len(self.layers)
-
-        j = self.args.sliding_window_pattern
-        full_mask = create_attention_mask(h, cache[j - 1])
-        swa_mask = create_attention_mask(h, cache[0], window_size=self.window_size)
-
-        for i, (layer, c) in enumerate(zip(self.layers, cache)):
-            is_global = (
-                i % self.args.sliding_window_pattern
-                == self.args.sliding_window_pattern - 1
-            )
-
-            mask = full_mask if is_global else swa_mask
-
-            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.model = CohereModel(args)
-        self.args = args
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        cache=None,
-    ):
-        out = self.model(inputs, 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,13 +1,11 @@
-# Copyright © 2023-2024 Apple Inc.
-
 from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Optional, Tuple
 
 import mlx.core as mx
 import mlx.nn as nn
 import numpy as np
 
-from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .base import BaseModelArgs
 
 
 @dataclass
@@ -49,7 +47,7 @@ class Attention(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = None,
     ) -> mx.array:
 
         qkv = self.Wqkv(x)
@@ -74,8 +72,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = scaled_dot_product_attention(
-            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        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)
@@ -92,7 +90,7 @@ class NormAttnNorm(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = None,
     ) -> mx.array:
         h = self.attn(self.norm_1(x), mask=mask, cache=cache)
         x = h + x
@@ -105,9 +103,10 @@ 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 = nn.silu(self.w1(x)) * self.v1(x)
+        current_hidden_states = self.act_fn(self.w1(x)) * self.v1(x)
         current_hidden_states = self.w2(current_hidden_states)
         return current_hidden_states
 
@@ -178,7 +177,7 @@ class DecoderLayer(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = None,
     ) -> mx.array:
         r, h = self.norm_attn_norm(x, mask, cache)
         out = self.ffn(h) + r
@@ -200,11 +199,15 @@ class DBRX(nn.Module):
     ):
         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 cache is None:
             cache = [None] * len(self.blocks)
 
-        mask = create_attention_mask(h, cache[0])
-
         for layer, c in zip(self.blocks, cache):
             h = layer(h, mask, c)
 
@@ -248,3 +251,11 @@ 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

@@ -1,258 +0,0 @@
-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,
-    ) -> mx.array:
-        h = self.embed_tokens(x)
-
-        if cache is None:
-            cache = [None] * len(self.layers)
-
-        mask = create_attention_mask(h, cache[0])
-
-        for layer, c in zip(self.layers, cache):
-            h = layer(h, mask, c)
-
-        return self.norm(h)
-
-
-class Model(nn.Module):
-    def __init__(self, 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,
-    ):
-        out = self.model(inputs, cache)
-        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

@@ -1,434 +0,0 @@
-# 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 .pipeline import PipelineMixin
-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(PipelineMixin, 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.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-
-    def __call__(
-        self,
-        x: mx.array,
-        cache: Optional[Any] = None,
-    ) -> mx.array:
-        h = self.embed_tokens(x)
-
-        pipeline_rank = self.pipeline_rank
-        pipeline_size = self.pipeline_size
-
-        if cache is None:
-            cache = [None] * len(self.pipeline_layers)
-        mask = create_attention_mask(h, cache[0])
-
-        # Receive from the previous process in the pipeline
-        if pipeline_rank < pipeline_size - 1:
-            h = mx.distributed.recv_like(h, (pipeline_rank + 1))
-
-        for l, c in zip(self.pipeline_layers, cache):
-            h = l(h, mask, cache=c)
-
-        # Send to the next process in the pipeline
-        if pipeline_rank != 0:
-            h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
-            if cache[-1] is not None:
-                cache[-1].keys = mx.depends(cache[-1].keys, h)
-
-        # Broadcast h while keeping it in the graph
-        h = mx.distributed.all_gather(h)[: h.shape[0]]
-
-        return self.norm(h)
-
-
-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,
-    ):
-        out = self.model(inputs, cache)
-        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.pipeline_layers

mlx_lm/models/deepseek_v3.py

@@ -1,431 +0,0 @@
-# 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 .pipeline import PipelineMixin
-from .rope_utils import initialize_rope
-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: int = 1
-    topk_group: int = 1
-    num_experts_per_tok: int = 1
-    moe_layer_freq: int = 1
-    first_k_dense_replace: int = 0
-    max_position_embeddings: int = 2048
-    rms_norm_eps: float = 1e-6
-    rope_theta: float = 10000.0
-    rope_scaling: Dict = None
-    attention_bias: bool = False
-
-
-class 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)
-            if mscale_all_dim:
-                scaling_factor = self.config.rope_scaling["factor"]
-                if scaling_factor > 1:
-                    s = 0.1 * mscale_all_dim * math.log(scaling_factor) + 1.0
-                    self.scale = self.scale * s * s
-
-        self.rope = initialize_rope(
-            dims=self.qk_rope_head_dim,
-            base=self.rope_theta,
-            traditional=True,
-            max_position_embeddings=self.max_position_embeddings,
-            scaling_config=self.config.rope_scaling,
-        )
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-    ) -> mx.array:
-        B, L, D = x.shape
-
-        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,
-):
-
-    scores = mx.sigmoid(gates.astype(mx.float32))
-    orig_scores = scores
-    scores = scores + e_score_correction_bias
-    if n_group > 1:
-        scores = mx.unflatten(scores, axis=-1, shape=(n_group, -1))
-        group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True)
-        k = n_group - topk_group
-        group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :]
-        scores = mx.put_along_axis(
-            scores, mx.stop_gradient(group_idx), mx.array(0.0), axis=-2
-        )
-        scores = mx.flatten(scores, -2, -1)
-
-    k = top_k
-    inds = mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k]
-    scores = mx.take_along_axis(orig_scores, inds, axis=-1)
-    if top_k > 1 and norm_topk_prob:
-        denominator = scores.sum(axis=-1, keepdims=True)
-        scores = scores / denominator
-    scores = scores * routed_scaling_factor
-
-    return inds, scores
-
-
-class MoEGate(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.config = config
-        self.top_k = config.num_experts_per_tok
-        self.norm_topk_prob = config.norm_topk_prob
-        self.n_routed_experts = config.n_routed_experts
-        self.routed_scaling_factor = config.routed_scaling_factor
-        self.n_group = config.n_group
-        self.topk_group = config.topk_group
-        self.weight = mx.zeros((self.n_routed_experts, config.hidden_size))
-        self.e_score_correction_bias = mx.zeros((self.n_routed_experts,))
-        assert config.topk_method == "noaux_tc", "Unsupported topk method."
-
-    def __call__(self, x):
-        return group_expert_select(
-            x @ self.weight.T,
-            self.e_score_correction_bias,
-            self.top_k,
-            self.n_group,
-            self.topk_group,
-            self.routed_scaling_factor,
-            self.norm_topk_prob,
-        )
-
-
-class 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,
-        )
-
-        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(PipelineMixin, 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.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-
-    def __call__(
-        self,
-        x: mx.array,
-        cache: Optional[Any] = None,
-    ) -> mx.array:
-        h = self.embed_tokens(x)
-
-        pipeline_rank = self.pipeline_rank
-        pipeline_size = self.pipeline_size
-
-        if cache is None:
-            cache = [None] * len(self.pipeline_layers)
-        mask = create_attention_mask(h, cache[0])
-
-        # Receive from the previous process in the pipeline
-        if pipeline_rank < pipeline_size - 1:
-            h = mx.distributed.recv_like(h, (pipeline_rank + 1))
-
-        for l, c in zip(self.pipeline_layers, cache):
-            h = l(h, mask, cache=c)
-
-        # Send to the next process in the pipeline
-        if pipeline_rank != 0:
-            h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
-            if cache[-1] is not None:
-                cache[-1].keys = mx.depends(cache[-1].keys, h)
-
-        # Broadcast h while keeping it in the graph
-        h = mx.distributed.all_gather(h)[: h.shape[0]]
-
-        return self.norm(h)
-
-
-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,
-    ):
-        out = self.model(inputs, cache)
-        return self.lm_head(out)
-
-    def sanitize(self, weights):
-        def dequant(weight, scale_inv):
-            dtype = weight.dtype
-            bs = 128  # block size
-            m, n = weight.shape
-            pad_bottom = (-m) % bs
-            pad_side = (-n) % bs
-            weight = mx.pad(weight, ((0, pad_bottom), (0, pad_side)))
-            weight = weight.reshape(
-                ((m + pad_bottom) // bs, bs, (n + pad_side) // bs, bs)
-            )
-            weight = (weight * scale_inv[:, None, :, None]).reshape(
-                m + pad_bottom, n + pad_side
-            )
-            return weight[:m, :n].astype(dtype)
-
-        # Remap for int4
-        new_weights = {}
-        for k, v in weights.items():
-            if k.endswith("weight_shape"):
-                base = k.replace("weight_shape", "")
-                new_weights[base + "weight"] = weights[base + "weight_packed"].view(
-                    mx.uint32
-                )
-                s = weights[base + "weight_scale"]
-                new_weights[base + "scales"] = s
-                new_weights[base + "biases"] = -8 * s
-            elif not (k.endswith("weight_scale") or k.endswith("weight_packed")):
-                new_weights[k] = v
-        weights = new_weights
-
-        # Dequantize fp8
-        new_weights = {}
-        for k, v in weights.items():
-            if "weight_scale_inv" in k:
-                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.pipeline_layers
-
-    @property
-    def cast_predicate(self):
-        def predicate(k):
-            return "e_score_correction_bias" not in k
-
-        return predicate

mlx_lm/models/deepseek_v32.py

@@ -1,515 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-import math
-from dataclasses import dataclass
-from typing import Any, Dict, Optional
-
-import mlx.core as mx
-import mlx.nn as nn
-
-from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
-from .cache import CacheList, KVCache
-from .rope_utils import initialize_rope
-from .switch_layers import SwitchGLU
-
-
-@dataclass
-class ModelArgs(BaseModelArgs):
-    model_type: str = "deepseek_v32"
-    vocab_size: int = 102400
-    hidden_size: int = 4096
-    index_head_dim: int = 128
-    index_n_heads: int = 64
-    index_topk: int = 2048
-    intermediate_size: int = 11008
-    moe_intermediate_size: int = 1407
-    num_hidden_layers: int = 30
-    num_attention_heads: int = 32
-    num_key_value_heads: int = 32
-    n_shared_experts: Optional[int] = None
-    n_routed_experts: Optional[int] = None
-    routed_scaling_factor: float = 1.0
-    kv_lora_rank: int = 512
-    q_lora_rank: int = 1536
-    qk_rope_head_dim: int = 64
-    v_head_dim: int = 128
-    qk_nope_head_dim: int = 128
-    topk_method: str = "noaux_tc"
-    scoring_func: str = "sigmoid"
-    norm_topk_prob: bool = True
-    n_group: int = 1
-    topk_group: int = 1
-    num_experts_per_tok: int = 1
-    moe_layer_freq: int = 1
-    first_k_dense_replace: int = 0
-    max_position_embeddings: int = 2048
-    rms_norm_eps: float = 1e-6
-    rope_theta: float = 10000.0
-    rope_scaling: Dict = None
-    attention_bias: bool = False
-
-
-class Indexer(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.dim = args.hidden_size
-        self.n_heads = args.index_n_heads
-        self.head_dim = args.index_head_dim
-        self.rope_head_dim = args.qk_rope_head_dim
-        self.index_topk = args.index_topk
-        self.q_lora_rank = args.q_lora_rank
-        self.wq_b = nn.Linear(
-            self.q_lora_rank, self.n_heads * self.head_dim, bias=False
-        )
-        self.wk = nn.Linear(self.dim, self.head_dim, bias=False)
-        self.k_norm = nn.LayerNorm(self.head_dim)
-        self.weights_proj = nn.Linear(self.dim, self.n_heads, bias=False)
-        self.softmax_scale = self.head_dim**-0.5
-        self.rope = initialize_rope(
-            dims=args.qk_rope_head_dim,
-            base=args.rope_theta,
-            traditional=False,
-            max_position_embeddings=args.max_position_embeddings,
-            scaling_config=args.rope_scaling,
-        )
-
-    def __call__(
-        self,
-        x: mx.array,
-        qr: mx.array,
-        mask: Optional[mx.array],
-        cache: Optional[Any] = None,
-    ):
-        # Computes top_k indices for attention
-        b, s, _ = x.shape
-        q = self.wq_b(qr)
-        q = q.reshape(b, s, self.n_heads, self.head_dim).swapaxes(1, 2)
-        q_pe, q_nope = mx.split(q, [self.rope_head_dim], axis=-1)
-
-        offset = cache.offset if cache is not None else 0
-
-        q_pe = self.rope(q_pe, offset=offset)
-        q = mx.concatenate([q_pe, q_nope], axis=-1)
-
-        k = self.wk(x)
-        k = self.k_norm(k)
-        k = mx.reshape(k, (b, 1, s, self.head_dim))
-        k_pe, k_nope = mx.split(k, [self.rope_head_dim], axis=-1)
-        k_pe = self.rope(k_pe, offset=offset)
-        k = mx.concatenate([k_pe, k_nope], axis=-1)
-        if cache is not None:
-            k, _ = cache.update_and_fetch(k, mx.zeros([b, 1, s, 0]))
-        if k.shape[2] <= self.index_topk:
-            return None
-        scores = q @ k.swapaxes(-1, -2)
-        scores = mx.maximum(scores, 0)
-        weights = self.weights_proj(x) * (self.n_heads**-0.5 * self.softmax_scale)
-        weights = weights.swapaxes(-1, -2)[..., None]
-        scores = scores * weights
-        scores = scores.sum(axis=1)
-        if mask is not None:
-            scores = mx.where(mask, scores, -float("inf"))
-        return mx.argpartition(scores, kth=-self.index_topk, axis=-1)[
-            ..., -self.index_topk :
-        ]
-
-
-class DeepseekV32Attention(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.config = config
-        self.hidden_size = config.hidden_size
-        self.num_heads = config.num_attention_heads
-        self.max_position_embeddings = config.max_position_embeddings
-        self.rope_theta = config.rope_theta
-        self.q_lora_rank = config.q_lora_rank
-        self.qk_rope_head_dim = config.qk_rope_head_dim
-        self.kv_lora_rank = config.kv_lora_rank
-        self.v_head_dim = config.v_head_dim
-        self.qk_nope_head_dim = config.qk_nope_head_dim
-        self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim
-
-        self.scale = self.q_head_dim**-0.5
-
-        self.q_a_proj = nn.Linear(
-            self.hidden_size, self.q_lora_rank, bias=config.attention_bias
-        )
-        self.q_a_layernorm = nn.RMSNorm(self.q_lora_rank, eps=1e-6)
-        self.q_b_proj = nn.Linear(
-            self.q_lora_rank, self.num_heads * self.q_head_dim, bias=False
-        )
-
-        self.kv_a_proj_with_mqa = nn.Linear(
-            self.hidden_size,
-            self.kv_lora_rank + self.qk_rope_head_dim,
-            bias=config.attention_bias,
-        )
-        self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank, eps=1e-6)
-        self.kv_b_proj = nn.Linear(
-            self.kv_lora_rank,
-            self.num_heads
-            * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
-            bias=False,
-        )
-
-        self.o_proj = nn.Linear(
-            self.num_heads * self.v_head_dim,
-            self.hidden_size,
-            bias=config.attention_bias,
-        )
-
-        if self.config.rope_scaling is not None:
-            mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0)
-            if mscale_all_dim:
-                scaling_factor = self.config.rope_scaling["factor"]
-                if scaling_factor > 1:
-                    s = 0.1 * mscale_all_dim * math.log(scaling_factor) + 1.0
-                    self.scale = self.scale * s * s
-
-        self.indexer = Indexer(config)
-        self.rope = initialize_rope(
-            dims=self.qk_rope_head_dim,
-            base=self.rope_theta,
-            traditional=True,
-            max_position_embeddings=self.max_position_embeddings,
-            scaling_config=self.config.rope_scaling,
-        )
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-    ) -> mx.array:
-        B, L, D = x.shape
-
-        qr = self.q_a_layernorm(self.q_a_proj(x))
-        q = self.q_b_proj(qr)
-
-        q = q.reshape(B, L, self.num_heads, self.q_head_dim).transpose(0, 2, 1, 3)
-        q_nope, q_pe = mx.split(q, [self.qk_nope_head_dim], axis=-1)
-        compressed_kv = self.kv_a_proj_with_mqa(x)
-        compressed_kv, k_pe = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
-        k_pe = k_pe.reshape(B, L, 1, self.qk_rope_head_dim).transpose(0, 2, 1, 3)
-        kv = self.kv_b_proj(self.kv_a_layernorm(compressed_kv))
-        kv = kv.reshape(B, L, self.num_heads, -1).transpose(0, 2, 1, 3)
-
-        k_nope, values = mx.split(kv, [self.qk_nope_head_dim], axis=-1)
-
-        if cache is not None:
-            q_pe = self.rope(q_pe, cache[0].offset)
-            k_pe = self.rope(k_pe, cache[0].offset)
-            k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
-            keys, values = cache[0].update_and_fetch(
-                mx.concatenate([k_nope, k_pe], axis=-1), values
-            )
-        else:
-            cache = [None] * 2
-            q_pe = self.rope(q_pe)
-            k_pe = self.rope(k_pe)
-            k_pe = mx.repeat(k_pe, self.num_heads, axis=1)
-            keys = mx.concatenate([k_nope, k_pe], axis=-1)
-
-        queries = mx.concatenate([q_nope, q_pe], axis=-1)
-        topk_indices = self.indexer(x, qr, mask, cache=cache[1])
-        if topk_indices is not None:
-            k_seq = keys.shape[2]
-            sparse_mask = mx.zeros((B, L, k_seq), dtype=mx.bool_)
-            sparse_mask = mx.put_along_axis(
-                sparse_mask, topk_indices, mx.array(True), axis=-1
-            )
-            sparse_mask = sparse_mask[:, None, :, :]
-            if mask is not None:
-                sparse_mask = sparse_mask & mask
-            mask = sparse_mask
-        output = scaled_dot_product_attention(
-            queries, keys, values, cache=cache[0], scale=self.scale, mask=mask
-        )
-        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
-        return self.o_proj(output)
-
-
-class DeepseekV32MLP(nn.Module):
-    def __init__(
-        self, config: ModelArgs, hidden_size: int = None, intermediate_size: int = None
-    ):
-        super().__init__()
-        self.config = config
-        self.hidden_size = config.hidden_size if hidden_size is None else hidden_size
-        self.intermediate_size = (
-            config.intermediate_size if intermediate_size is None else intermediate_size
-        )
-
-        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
-
-    def __call__(self, x):
-        down_proj = self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
-        return down_proj
-
-
-@mx.compile
-def group_expert_select(
-    gates,
-    e_score_correction_bias,
-    top_k,
-    n_group,
-    topk_group,
-    routed_scaling_factor,
-    norm_topk_prob,
-):
-
-    scores = mx.sigmoid(gates.astype(mx.float32))
-    orig_scores = scores
-    scores = scores + e_score_correction_bias
-    if n_group > 1:
-        scores = mx.unflatten(scores, axis=-1, shape=(n_group, -1))
-        group_scores = mx.topk(scores, 2, axis=-1).sum(axis=-1, keepdims=True)
-        k = n_group - topk_group
-        group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :]
-        scores = mx.put_along_axis(
-            scores, mx.stop_gradient(group_idx), mx.array(0.0), axis=-2
-        )
-        scores = mx.flatten(scores, -2, -1)
-
-    k = top_k
-    inds = mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k]
-    scores = mx.take_along_axis(orig_scores, inds, axis=-1)
-    if top_k > 1 and norm_topk_prob:
-        denominator = scores.sum(axis=-1, keepdims=True)
-        scores = scores / denominator
-    scores = scores * routed_scaling_factor
-
-    return inds, scores
-
-
-class MoEGate(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.config = config
-        self.top_k = config.num_experts_per_tok
-        self.norm_topk_prob = config.norm_topk_prob
-        self.n_routed_experts = config.n_routed_experts
-        self.routed_scaling_factor = config.routed_scaling_factor
-        self.n_group = config.n_group
-        self.topk_group = config.topk_group
-        self.weight = mx.zeros((self.n_routed_experts, config.hidden_size))
-        self.e_score_correction_bias = mx.zeros((self.n_routed_experts,))
-        assert config.topk_method == "noaux_tc", "Unsupported topk method."
-
-    def __call__(self, x):
-        return group_expert_select(
-            x @ self.weight.T,
-            self.e_score_correction_bias,
-            self.top_k,
-            self.n_group,
-            self.topk_group,
-            self.routed_scaling_factor,
-            self.norm_topk_prob,
-        )
-
-
-class DeepseekV32MoE(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.config = config
-        self.num_experts_per_tok = config.num_experts_per_tok
-        self.switch_mlp = SwitchGLU(
-            config.hidden_size,
-            config.moe_intermediate_size,
-            config.n_routed_experts,
-        )
-
-        self.gate = MoEGate(config)
-        if config.n_shared_experts is not None:
-            intermediate_size = config.moe_intermediate_size * config.n_shared_experts
-            self.shared_experts = DeepseekV32MLP(
-                config=config, intermediate_size=intermediate_size
-            )
-
-    def __call__(self, x):
-        inds, scores = self.gate(x)
-        y = self.switch_mlp(x, inds)
-        y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
-        if self.config.n_shared_experts is not None:
-            y = y + self.shared_experts(x)
-
-        return y
-
-
-class DeepseekV32DecoderLayer(nn.Module):
-    def __init__(self, config: ModelArgs, layer_idx: int):
-        super().__init__()
-        self.self_attn = DeepseekV32Attention(config)
-        self.mlp = (
-            DeepseekV32MoE(config)
-            if (
-                config.n_routed_experts is not None
-                and layer_idx >= config.first_k_dense_replace
-                and layer_idx % config.moe_layer_freq == 0
-            )
-            else DeepseekV32MLP(config)
-        )
-        self.input_layernorm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.post_attention_layernorm = nn.RMSNorm(
-            config.hidden_size, eps=config.rms_norm_eps
-        )
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-    ) -> mx.array:
-        r = self.self_attn(self.input_layernorm(x), mask, cache)
-        h = x + r
-        r = self.mlp(self.post_attention_layernorm(h))
-        return h + r
-
-
-class DeepseekV32Model(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.vocab_size = config.vocab_size
-        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
-        self.layers = [
-            DeepseekV32DecoderLayer(config, idx)
-            for idx in range(config.num_hidden_layers)
-        ]
-        self.start_idx = 0
-        self.end_idx = len(self.layers)
-        self.num_layers = self.end_idx
-
-        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.pipeline_rank = 0
-        self.pipeline_size = 1
-
-    def pipeline(self, group):
-        # Split layers in reverse so rank=0 gets the last layers and
-        # rank=pipeline_size-1 gets the first
-        self.pipeline_rank = group.rank()
-        self.pipeline_size = group.size()
-        layers_per_rank = len(self.layers) // self.pipeline_size
-        extra = len(self.layers) - layers_per_rank * self.pipeline_size
-        if self.pipeline_rank < extra:
-            layers_per_rank += 1
-        self.start_idx = (self.pipeline_size - self.pipeline_rank - 1) * layers_per_rank
-        self.end_idx = self.start_idx + layers_per_rank
-        self.layers = self.layers[: self.end_idx]
-        self.layers[: self.start_idx] = [None] * self.start_idx
-        self.num_layers = len(self.layers) - self.start_idx
-
-    def __call__(
-        self,
-        x: mx.array,
-        cache: Optional[Any] = None,
-    ) -> mx.array:
-        h = self.embed_tokens(x)
-
-        pipeline_rank = self.pipeline_rank
-        pipeline_size = self.pipeline_size
-
-        if cache is None:
-            cache = [None] * self.num_layers
-        mask = create_attention_mask(
-            h, cache[0][0] if cache[0] else None, return_array=True
-        )
-
-        # Receive from the previous process in the pipeline
-
-        if pipeline_rank < pipeline_size - 1:
-            h = mx.distributed.recv_like(h, (pipeline_rank + 1))
-
-        for i in range(self.num_layers):
-            h = self.layers[self.start_idx + i](h, mask, cache[i])
-
-        # Send to the next process in the pipeline
-        if pipeline_rank != 0:
-            h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
-            if cache[-1] is not None:
-                cache[-1].keys = mx.depends(cache[-1].keys, h)
-
-        # Broadcast h while keeping it in the graph
-        h = mx.distributed.all_gather(h)[: h.shape[0]]
-
-        return self.norm(h)
-
-
-class Model(nn.Module):
-    def __init__(self, config: ModelArgs):
-        super().__init__()
-        self.args = config
-        self.model_type = config.model_type
-        self.model = DeepseekV32Model(config)
-        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        cache: Optional[Any] = None,
-    ):
-        out = self.model(inputs, cache)
-        return self.lm_head(out)
-
-    def sanitize(self, weights):
-        def dequant(weight, scale_inv):
-            dtype = weight.dtype
-            bs = 128  # block size
-            m, n = weight.shape
-            pad_bottom = (-m) % bs
-            pad_side = (-n) % bs
-            weight = mx.pad(weight, ((0, pad_bottom), (0, pad_side)))
-            weight = weight.reshape(
-                ((m + pad_bottom) // bs, bs, (n + pad_side) // bs, bs)
-            )
-            weight = (weight * scale_inv[:, None, :, None]).reshape(
-                m + pad_bottom, n + pad_side
-            )
-            return weight[:m, :n].astype(dtype)
-
-        # Dequantize
-        new_weights = {}
-        for k, v in weights.items():
-            if "weight_scale_inv" in k:
-                scale_inv = v
-                wk = k.replace("_scale_inv", "")
-                weight = weights[wk]
-                weight = dequant(weight, scale_inv)
-                new_weights[wk] = weight
-            elif k not in new_weights:
-                new_weights[k] = v
-        weights = new_weights
-
-        # Stack experts
-        for l in range(self.args.num_hidden_layers):
-            prefix = f"model.layers.{l}"
-            for n, m in [("w1", "gate_proj"), ("w2", "down_proj"), ("w3", "up_proj")]:
-                for k in ["weight", "scales", "biases"]:
-                    if f"{prefix}.mlp.experts.0.{m}.{k}" in weights:
-                        to_join = [
-                            weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}")
-                            for e in range(self.args.n_routed_experts)
-                        ]
-                        weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
-
-        # Remove multi-token prediction layer and any unused precomputed rotary freqs
-        return {
-            k: v
-            for k, v in weights.items()
-            if not k.startswith("model.layers.61") and "rotary_emb.inv_freq" not in k
-        }
-
-    @property
-    def layers(self):
-        return self.model.layers[self.model.start_idx : self.model.end_idx]
-
-    @property
-    def cast_predicate(self):
-        def predicate(k):
-            return "e_score_correction_bias" not in k
-
-        return predicate
-
-    def make_cache(self):
-        return [CacheList(KVCache(), KVCache()) for _ in self.layers]

mlx_lm/models/dots1.py

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

mlx_lm/models/ernie4_5.py

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

mlx_lm/models/ernie4_5_moe.py

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

mlx_lm/models/exaone.py

@@ -1,164 +0,0 @@
-# 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,
-        cache=None,
-    ):
-        h = self.wte(inputs)
-
-        if cache is None:
-            cache = [None] * len(self.h)
-
-        mask = create_attention_mask(h, cache[0])
-
-        for layer, c in zip(self.h, cache):
-            h = layer(h, mask, cache=c)
-
-        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,
-        cache=None,
-    ):
-        out = self.transformer(inputs, 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/exaone4.py

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

mlx_lm/models/falcon_h1.py

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

mlx_lm/models/gated_delta.py

@@ -1,284 +0,0 @@
-from functools import partial
-from typing import Optional, Tuple
-
-import mlx.core as mx
-import mlx.nn as nn
-
-
-@partial(mx.compile, shapeless=True)
-def compute_g(A_log, a, dt_bias):
-    return mx.exp(-mx.exp(A_log.astype(mx.float32)) * nn.softplus(a + dt_bias)).astype(
-        A_log.dtype
-    )
-
-
-def _make_gated_delta_kernel(has_mask=False, vectorized=False):
-    if not mx.metal.is_available():
-        return None
-    mask_source = "mask[b_idx * T + t]" if has_mask else "true"
-
-    # Configure g indexing based on whether gating is vectorized
-    if vectorized:
-        g_comment = "// g: [B, T, Hv, Dk]"
-        g_setup = "auto g_ = g + (b_idx * T * Hv + hv_idx) * Dk;"
-        g_access = "g_[s_idx]"
-        g_advance = "g_ += Hv * Dk;"
-    else:
-        g_comment = "// g: [B, T, Hv]"
-        g_setup = "auto g_ = g + b_idx * T * Hv;"
-        g_access = "g_[hv_idx]"
-        g_advance = "g_ += Hv;"
-
-    source = f"""
-        auto n = thread_position_in_grid.z;
-        auto b_idx = n / Hv;
-        auto hv_idx = n % Hv;
-        auto hk_idx = hv_idx / (Hv / Hk);
-        constexpr int n_per_t = Dk / 32;
-
-        // q, k: [B, T, Hk, Dk]
-        auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk;
-        auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk;
-
-        // v, y: [B, T, Hv, Dv]
-        auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv;
-        y += b_idx * T * Hv * Dv + hv_idx * Dv;
-
-        auto dk_idx = thread_position_in_threadgroup.x;
-        auto dv_idx = thread_position_in_grid.y;
-
-        // state_in, state_out: [B, Hv, Dv, Dk]
-        auto i_state = state_in + (n * Dv + dv_idx) * Dk;
-        auto o_state = state_out + (n * Dv + dv_idx) * Dk;
-
-        float state[n_per_t];
-        for (int i = 0; i < n_per_t; ++i) {{
-          auto s_idx = n_per_t * dk_idx + i;
-          state[i] = static_cast<float>(i_state[s_idx]);
-        }}
-
-        {g_comment}
-        {g_setup}
-        auto beta_ = beta + b_idx * T * Hv;
-
-        for (int t = 0; t < T; ++t) {{
-          if ({mask_source}) {{
-            float kv_mem = 0.0f;
-            for (int i = 0; i < n_per_t; ++i) {{
-              auto s_idx = n_per_t * dk_idx + i;
-              state[i] = state[i] * {g_access};
-              kv_mem += state[i] * k_[s_idx];
-            }}
-            kv_mem = simd_sum(kv_mem);
-
-            auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx];
-
-            float out = 0.0f;
-            for (int i = 0; i < n_per_t; ++i) {{
-              auto s_idx = n_per_t * dk_idx + i;
-              state[i] = state[i] + k_[s_idx] * delta;
-              out += state[i] * q_[s_idx];
-            }}
-            out = simd_sum(out);
-            if (thread_index_in_simdgroup == 0) {{
-              y[dv_idx] = static_cast<InT>(out);
-            }}
-          }}
-          // Increment data pointers to next time step
-          q_ += Hk * Dk;
-          k_ += Hk * Dk;
-          v_ += Hv * Dv;
-          y += Hv * Dv;
-          {g_advance}
-          beta_ += Hv;
-        }}
-        for (int i = 0; i < n_per_t; ++i) {{
-          auto s_idx = n_per_t * dk_idx + i;
-          o_state[s_idx] = static_cast<InT>(state[i]);
-        }}
-    """
-    inputs = ["q", "k", "v", "g", "beta", "state_in", "T"]
-    if has_mask:
-        inputs.append("mask")
-
-    suffix = ""
-    if vectorized:
-        suffix += "_vec"
-    if has_mask:
-        suffix += "_mask"
-
-    return mx.fast.metal_kernel(
-        name=f"gated_delta_step{suffix}",
-        input_names=inputs,
-        output_names=["y", "state_out"],
-        source=source,
-    )
-
-
-_gated_delta_kernel = _make_gated_delta_kernel(has_mask=False, vectorized=False)
-_gated_delta_kernel_masked = _make_gated_delta_kernel(has_mask=True, vectorized=False)
-_gated_delta_kernel_vec = _make_gated_delta_kernel(has_mask=False, vectorized=True)
-_gated_delta_kernel_vec_masked = _make_gated_delta_kernel(
-    has_mask=True, vectorized=True
-)
-
-
-@mx.compile
-def _gated_delta_step_ops(
-    q: mx.array,
-    k: mx.array,
-    v: mx.array,
-    g: mx.array,
-    beta: mx.array,
-    state: mx.array,
-    mask: Optional[mx.array] = None,
-) -> Tuple[mx.array, mx.array]:
-    """
-    Ops-based reference implementation for a single recurrent step.
-
-    Shapes:
-      - q, k: [B, H, Dk]
-      - v: [B, H, Dv]
-      - g: [B, H] or [B, H, Dk]
-      - beta: [B, H]
-      - state: [B, H, Dv, Dk]
-    Returns:
-      - y: [B, H, Dv]
-      - new_state: [B, H, Dv, Dk]
-    """
-
-    # Decay
-    old_state = state
-    if g.ndim == 2:
-        decay = g[..., None, None]
-    elif g.ndim == 3:
-        decay = g[..., None, :]
-    else:
-        raise ValueError(f"Unsupported gating shape {g.shape}")
-    state = state * decay
-    kv_mem = (state * k[..., None, :]).sum(axis=-1)  # [B, H, Dv]
-    delta = (v - kv_mem) * beta[..., None]  # [B, H, Dv]
-    state = state + k[..., None, :] * delta[..., None]
-    # Output projection along key dim with q
-    y = (state * q[..., None, :]).sum(axis=-1)  # [B, H, Dv]
-    if mask is not None:
-        if mask.ndim == 2:
-            mask = mx.expand_dims(mask, axes=(2, 3))
-        elif mask.ndim == 3:
-            mask = mx.expand_dims(mask, axis=-1)
-        state = mx.where(mask, state, old_state)
-    return y, state
-
-
-def gated_delta_kernel(
-    q: mx.array,
-    k: mx.array,
-    v: mx.array,
-    g: mx.array,
-    beta: mx.array,
-    state: mx.array,
-    mask: Optional[mx.array] = None,
-) -> Tuple[mx.array, mx.array]:
-    B, T, Hk, Dk = k.shape
-    Hv, Dv = v.shape[2:]
-    input_type = q.dtype
-    if g.ndim == 4:
-        kernel = _gated_delta_kernel_vec
-        inputs = [q, k, v, g, beta, state, T]
-        if mask is not None:
-            kernel = _gated_delta_kernel_vec_masked
-            inputs.append(mask)
-    else:
-        kernel = _gated_delta_kernel
-        inputs = [q, k, v, g, beta, state, T]
-        if mask is not None:
-            kernel = _gated_delta_kernel_masked
-            inputs.append(mask)
-
-    return kernel(
-        inputs=inputs,
-        template=[
-            ("InT", input_type),
-            ("Dk", Dk),
-            ("Dv", Dv),
-            ("Hk", Hk),
-            ("Hv", Hv),
-        ],
-        grid=(32, Dv, B * Hv),
-        threadgroup=(32, 4, 1),
-        output_shapes=[(B, T, Hv, Dv), state.shape],
-        output_dtypes=[input_type, input_type],
-    )
-
-
-def gated_delta_ops(
-    q: mx.array,
-    k: mx.array,
-    v: mx.array,
-    g: mx.array,
-    beta: mx.array,
-    state: Optional[mx.array] = None,
-    mask: Optional[mx.array] = None,
-) -> Tuple[mx.array, mx.array]:
-    """
-    Ops-based reference implementation for prompt prefill (sequential loop).
-    Supports both scalar and vectorized gating.
-
-    Shapes:
-      - q, k: [B, T, Hk, Dk]
-      - v: [B, T, Hv, Dv]
-      - g: [B, T, Hv] (scalar) or [B, T, Hv, Dk] (vectorized)
-      - beta: [B, T, Hv]
-      - state: [B, Hv, Dv, Dk]
-    Returns:
-      - y: [B, T, Hv, Dv]
-      - state: [B, Hv, Dv, Dk]
-    """
-    B, T, Hk, Dk = q.shape
-    Hv, Dv = v.shape[-2:]
-    if state is None:
-        state = mx.zeros((B, Hv, Dv, Dk), dtype=q.dtype)
-
-    if (repeat_factor := Hv // Hk) > 1:
-        q = mx.repeat(q, repeat_factor, -2)
-        k = mx.repeat(k, repeat_factor, -2)
-
-    ys = []
-    for t in range(T):
-        y, state = _gated_delta_step_ops(
-            q[:, t],
-            k[:, t],
-            v[:, t],
-            g[:, t],
-            beta[:, t],
-            state,
-            None if mask is None else mask[:, t],
-        )
-        ys.append(y)
-    y = mx.stack(ys, axis=1)
-    return y, state
-
-
-def gated_delta_update(
-    q: mx.array,
-    k: mx.array,
-    v: mx.array,
-    a: mx.array,
-    b: mx.array,
-    A_log: mx.array,
-    dt_bias: mx.array,
-    state: Optional[mx.array] = None,
-    mask: Optional[mx.array] = None,
-    use_kernel: bool = True,
-) -> Tuple[mx.array, mx.array]:
-
-    beta = mx.sigmoid(b)
-    g = compute_g(A_log, a, dt_bias)
-    if state is None:
-        B, _, Hk, Dk = q.shape
-        Hv, Dv = v.shape[-2:]
-        state = mx.zeros((B, Hv, Dv, Dk), dtype=q.dtype)
-
-    if not use_kernel or mx.default_device() != mx.gpu or not mx.metal.is_available():
-        return gated_delta_ops(q, k, v, g, beta, state, mask)
-    return gated_delta_kernel(q, k, v, g, beta, state, mask)

mlx_lm/models/gemma.py

@@ -1,12 +1,10 @@
-# Copyright © 2023-2024 Apple Inc.
-
 from dataclasses import dataclass
-from typing import Any, Optional
+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 .base import BaseModelArgs
 
 
 @dataclass
@@ -60,7 +58,7 @@ class Attention(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = None,
     ) -> mx.array:
         B, L, D = x.shape
 
@@ -79,8 +77,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = scaled_dot_product_attention(
-            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        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)
@@ -113,7 +111,7 @@ class TransformerBlock(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = None,
     ) -> mx.array:
         r = self.self_attn(self.input_layernorm(x), mask, cache)
         h = x + r
@@ -143,11 +141,14 @@ class GemmaModel(nn.Module):
         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 cache is None:
             cache = [None] * len(self.layers)
 
-        mask = create_attention_mask(h, cache[0])
-
         for layer, c in zip(self.layers, cache):
             h = layer(h, mask, c)
 
@@ -173,3 +174,11 @@ class Model(nn.Module):
     @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

@@ -1,205 +0,0 @@
-# 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,
-        cache=None,
-    ):
-        h = self.embed_tokens(inputs)
-        h = h * (self.args.hidden_size**0.5)
-
-        if cache is None:
-            cache = [None] * len(self.layers)
-
-        mask = create_attention_mask(h, cache[0], return_array=True)
-
-        for layer, c in zip(self.layers, cache):
-            h = layer(h, mask, c)
-
-        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,
-        cache=None,
-    ):
-        out = self.model(inputs, cache)
-        out = self.model.embed_tokens.as_linear(out)
-        out = mx.tanh(out / self.final_logit_softcapping)
-        out = out * self.final_logit_softcapping
-        return out
-
-    @property
-    def layers(self):
-        return self.model.layers

mlx_lm/models/gemma3.py

@@ -1,63 +0,0 @@
-# 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,
-        input_embeddings: Optional[mx.array] = None,
-    ):
-        return self.language_model(
-            inputs, cache=cache, input_embeddings=input_embeddings
-        )
-
-    def sanitize(self, weights):
-        weights = tree_unflatten(list(weights.items()))
-        weights.pop("vision_tower", None)
-        weights.pop("multi_modal_projector", None)
-        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

@@ -1,250 +0,0 @@
-# Copyright © 2025 Apple Inc.
-
-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 .cache import KVCache, RotatingKVCache
-from .rope_utils import initialize_rope
-
-
-@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_theta: float = 1_000_000.0
-    rope_local_base_freq: float = 10_000.0
-    query_pre_attn_scalar: float = 256
-    sliding_window: int = 512
-    sliding_window_pattern: int = 6
-    max_position_embeddings: int = 32768
-    rope_scaling: Dict = None
-
-
-class Attention(nn.Module):
-    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 = initialize_rope(
-            dims=head_dim,
-            base=(args.rope_local_base_freq if self.is_sliding else args.rope_theta),
-            traditional=False,
-            max_position_embeddings=args.max_position_embeddings,
-            scaling_config=args.rope_scaling,
-        )
-
-    def __call__(
-        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
-        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.window_size = args.sliding_window
-        self.sliding_window_pattern = args.sliding_window_pattern
-        self.vocab_size = args.vocab_size
-        self.num_hidden_layers = args.num_hidden_layers
-        assert self.vocab_size > 0
-        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,
-        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)
-
-        global_mask = create_attention_mask(h, cache[self.sliding_window_pattern - 1])
-
-        if self.sliding_window_pattern > 1:
-            sliding_window_mask = create_attention_mask(
-                h,
-                cache[0],
-                window_size=self.window_size,
-            )
-        else:
-            sliding_window_mask = None
-        for i, (layer, c) in enumerate(zip(self.layers, cache)):
-            is_global = (
-                i % self.sliding_window_pattern == self.sliding_window_pattern - 1
-            )
-            mask = global_mask if is_global else sliding_window_mask
-            h = layer(h, mask, c)
-
-        return self.norm(h)
-
-
-class Model(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.model_type = args.model_type
-        self.model = Gemma3Model(args)
-        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
-        self.tie_word_embeddings = False
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        cache=None,
-        input_embeddings: Optional[mx.array] = None,
-    ):
-        out = self.model(inputs, cache, input_embeddings)
-        if self.tie_word_embeddings:
-            out = self.model.embed_tokens.as_linear(out)
-        else:
-            out = self.lm_head(out)
-        return out
-
-    def sanitize(self, weights):
-        if "lm_head.weight" not in weights:
-            self.tie_word_embeddings = True
-            self.pop("lm_head")
-        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))
-        return caches

mlx_lm/models/gemma3n.py

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

mlx_lm/models/glm.py

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

mlx_lm/models/glm4.py

@@ -1,180 +0,0 @@
-# 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,
-        cache: Optional[Any] = None,
-    ):
-        h = self.embed_tokens(inputs)
-
-        if cache is None:
-            cache = [None] * len(self.layers)
-
-        mask = create_attention_mask(h, cache[0])
-
-        for layer, c in zip(self.layers, cache):
-            h = layer(h, mask, cache=c)
-
-        return self.norm(h)
-
-
-class Model(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.model_type = args.model_type
-        self.model = Glm4Model(args)
-        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        cache: Optional[Any] = None,
-    ):
-        out = self.model(inputs, cache)
-        return self.lm_head(out)
-
-    @property
-    def layers(self):
-        return self.model.layers

mlx_lm/models/glm4_moe.py

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

mlx_lm/models/gpt2.py

@@ -1,12 +1,11 @@
-# Copyright © 2023 - 2024 Apple Inc.
-
 from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Dict, Optional, Tuple, Union
 
 import mlx.core as mx
 import mlx.nn as nn
+import numpy as np
 
-from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .base import BaseModelArgs, create_additive_causal_mask
 
 
 @dataclass
@@ -45,7 +44,7 @@ class Attention(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = None,
     ) -> mx.array:
         B, L, D = x.shape
 
@@ -60,8 +59,8 @@ class Attention(nn.Module):
         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 = 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)
@@ -99,7 +98,7 @@ class TransformerBlock(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = None,
     ) -> mx.array:
         r = self.attn(self.ln_1(x), mask, cache)
         h = x + r
@@ -131,20 +130,20 @@ class GPT2Model(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 cache is None:
             cache = [None] * len(self.h)
 
-        offset = 0
-        if cache[0] is not None:
-            offset = cache[0].offset
-
-        offset = mx.array(offset)
-        position_ids = mx.arange(L) + offset[..., None]
-
-        hidden_states += self.wpe(position_ids)
-
-        mask = create_attention_mask(hidden_states, cache[0])
-
         for layer, c in zip(self.h, cache):
             hidden_states = layer(hidden_states, mask, cache=c)
 
@@ -198,3 +197,11 @@ 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,13 +1,11 @@
-# Copyright © 2023-2024 Apple Inc.
-
 from dataclasses import dataclass
-from typing import Any, Optional
+from typing import Dict, Optional, Tuple, Union
 
 import mlx.core as mx
 import mlx.nn as nn
 import numpy as np
 
-from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .base import BaseModelArgs, create_additive_causal_mask
 
 
 @dataclass
@@ -57,7 +55,7 @@ class Attention(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = None,
     ) -> mx.array:
         B, L, D = x.shape
 
@@ -74,8 +72,8 @@ class Attention(nn.Module):
         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 = 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.c_proj(output)
@@ -114,7 +112,7 @@ class TransformerBlock(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = None,
     ) -> mx.array:
         r = self.attn(self.ln_1(x), mask, cache)
         h = x + r
@@ -143,15 +141,19 @@ 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 cache is None:
             cache = [None] * len(self.h)
-            position_ids = mx.array(np.arange(L))
-        else:
-            position_ids = mx.array(np.arange(cache[0].offset, cache[0].offset + L))
-
-        mask = create_attention_mask(hidden_states, cache[0])
-
-        hidden_states += self.wpe(position_ids)
 
         for layer, c in zip(self.h, cache):
             hidden_states = layer(hidden_states, mask, cache=c)
@@ -183,3 +185,11 @@ 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

@@ -1,225 +0,0 @@
-# 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
-    use_parallel_residual: bool = True
-    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.use_parallel_residual = args.use_parallel_residual
-        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:
-        if self.use_parallel_residual:
-            residual = x
-            # Run attention and feedforward network in parallel.
-            attn = self.attention(self.input_layernorm(x), mask, cache)
-            ffn = self.mlp(self.post_attention_layernorm(x))
-            out = attn + ffn + residual
-            return out
-        else:
-            # Run attention and feedforward network sequentially.
-            attn_output = self.attention(self.input_layernorm(x), mask, cache)
-            x = x + attn_output
-            ffn_output = self.mlp(self.post_attention_layernorm(x))
-            x = x + ffn_output
-            return x
-
-
-class GPTNeoXModel(nn.Module):
-    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,
-        cache=None,
-    ):
-        _, L = inputs.shape
-
-        hidden_states = self.embed_in(inputs)
-
-        if cache is None:
-            cache = [None] * len(self.h)
-
-        mask = create_attention_mask(hidden_states, cache[0])
-
-        for layer, c in zip(self.h, cache):
-            hidden_states = layer(hidden_states, mask, cache=c)
-
-        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,
-        cache=None,
-    ):
-        out = self.model(inputs, 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/gpt_oss.py

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

mlx_lm/models/granite.py

@@ -1,192 +0,0 @@
-# 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,
-        cache=None,
-    ):
-        h = self.embed_tokens(inputs) * self.embedding_multiplier
-
-        if cache is None:
-            cache = [None] * len(self.layers)
-
-        mask = create_attention_mask(h, cache[0])
-
-        for layer, c in zip(self.layers, cache):
-            h = layer(h, mask, cache=c)
-
-        return self.norm(h)
-
-
-class Model(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.model_type = args.model_type
-        self.model = 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,
-        cache=None,
-    ):
-        out = self.model(inputs, cache)
-        if self.args.tie_word_embeddings:
-            out = self.model.embed_tokens.as_linear(out)
-        else:
-            out = self.lm_head(out)
-        return out / self.logits_scaling
-
-    @property
-    def layers(self):
-        return self.model.layers

mlx_lm/models/granitemoe.py

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

mlx_lm/models/granitemoehybrid.py

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

mlx_lm/models/helium.py

@@ -1,182 +0,0 @@
-# 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,
-        cache=None,
-    ) -> mx.array:
-        h = self.embed_tokens(inputs)
-
-        if cache is None:
-            cache = [None] * len(self.layers)
-
-        mask = create_attention_mask(h, cache[0])
-
-        for layer, c in zip(self.layers, cache):
-            h = layer(h, mask, c)
-
-        return self.norm(h)
-
-
-class Model(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.model_type = args.model_type
-
-        self.model = 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,
-        cache=None,
-    ) -> mx.array:
-        out = self.model(inputs, cache)
-        if self.args.tie_word_embeddings:
-            out = self.model.embed_tokens.as_linear(out)
-        else:
-            out = self.lm_head(out)
-        return out
-
-    @property
-    def layers(self):
-        return self.model.layers

mlx_lm/models/hunyuan.py

@@ -1,333 +0,0 @@
-# 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,
-        cache=None,
-    ):
-        h = self.embed_tokens(inputs)
-
-        if cache is None:
-            cache = [None] * len(self.layers)
-
-        mask = create_attention_mask(h, cache[0])
-        for i, (layer, c) in enumerate(zip(self.layers, cache)):
-            if (not self.args.use_cla) or i % self.args.cla_share_factor == 0:
-                shared_kv_states = None
-            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,
-        cache=None,
-    ):
-        out = self.model(inputs, 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/hunyuan_v1_dense.py

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

mlx_lm/models/internlm2.py

@@ -1,12 +1,10 @@
-# Copyright © 2023-2024 Apple Inc.
-
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Union
+from typing import 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 .base import BaseModelArgs
 
 
 @dataclass
@@ -19,7 +17,6 @@ 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
@@ -35,50 +32,8 @@ 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 ["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,
-        )
+            if self.rope_scaling["type"] != "linear":
+                raise ValueError("rope_scaling 'type' currently only supports 'linear'")
 
 
 class Attention(nn.Module):
@@ -101,12 +56,10 @@ 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 2.0
+            else 1
         )
-
-        self.rope = DynamicNTKScalingRoPE(
+        self.rope = nn.RoPE(
             head_dim,
-            max_position_embeddings=args.max_position_embeddings,
             traditional=args.rope_traditional,
             base=args.rope_theta,
             scale=rope_scale,
@@ -116,7 +69,7 @@ class Attention(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = None,
     ) -> mx.array:
         B, L, D = x.shape
 
@@ -141,8 +94,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = scaled_dot_product_attention(
-            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        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.wo(output)
@@ -171,7 +124,7 @@ class TransformerBlock(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = None,
     ) -> mx.array:
         r = self.attention(self.attention_norm(x), mask, cache)
         h = x + r
@@ -197,10 +150,14 @@ class InternLM2Model(nn.Module):
     ):
         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 cache is None:
             cache = [None] * len(self.layers)
 
-        mask = create_attention_mask(h, cache[0])
         for layer, c in zip(self.layers, cache):
             h = layer(h, mask, cache=c)
 
@@ -228,10 +185,14 @@ class Model(nn.Module):
             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

@@ -1,237 +0,0 @@
-# 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,
-        cache=None,
-    ):
-        h = self.embed_tokens(inputs)
-
-        if cache is None:
-            cache = [None] * len(self.layers)
-
-        mask = create_attention_mask(h, cache[0])
-        for layer, c in zip(self.layers, cache):
-            h = layer(h, mask, cache=c)
-
-        return self.norm(h)
-
-
-class Model(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.model_type = args.model_type
-        self.model = 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,
-        cache=None,
-    ):
-        out = self.model(inputs, cache)
-        if self.args.tie_word_embeddings:
-            out = self.model.embed_tokens.as_linear(out)
-        else:
-            out = self.lm_head(out)
-        return out
-
-    def sanitize(self, weights):
-        # Remove unused precomputed rotary freqs
-        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/jamba.py

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

mlx_lm/models/kimi_linear.py

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

mlx_lm/models/kimi_vl.py

@@ -1,116 +0,0 @@
-# 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: int = 1
-    topk_group: int = 1
-    num_experts_per_tok: int = 1
-    moe_layer_freq: int = 1
-    first_k_dense_replace: int = 0
-    max_position_embeddings: int = 2048
-    rms_norm_eps: float = 1e-6
-    rope_theta: float = 10000.0
-    rope_scaling: Dict = None
-    attention_bias: bool = False
-
-
-@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,
-    ):
-        out = self.model(inputs, cache)
-        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,
-    ):
-        return self.language_model(inputs, cache)
-
-    def sanitize(self, weights):
-        def keep(key):
-            return (
-                "vision_tower" not in key
-                and "rotary_emb" not in key
-                and "multi_modal_projector" not in key
-            )
-
-        weights = {k: v for k, v in weights.items() if keep(k)}
-        # Stack experts
-        for l in range(self.args.text_config.num_hidden_layers):
-            prefix = f"language_model.model.layers.{l}"
-            for m in [("gate_proj"), ("down_proj"), ("up_proj")]:
-                for k in ["weight", "scales", "biases"]:
-                    if f"{prefix}.mlp.experts.0.{m}.{k}" in weights:
-                        to_join = [
-                            weights.pop(f"{prefix}.mlp.experts.{e}.{m}.{k}")
-                            for e in range(self.args.text_config.n_routed_experts)
-                        ]
-                        weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
-
-        return weights
-
-    @property
-    def layers(self):
-        return self.language_model.model.layers
-
-    @property
-    def cast_predicate(self):
-        def predicate(k):
-            return "e_score_correction_bias" not in k
-
-        return predicate

mlx_lm/models/lfm2-vl.py

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

mlx_lm/models/lfm2.py

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

mlx_lm/models/lfm2_moe.py

@@ -1,372 +0,0 @@
-# Copyright © 2025 Apple Inc.
-from dataclasses import dataclass
-from typing import Any, List, Optional
-
-import mlx.core as mx
-import mlx.nn as nn
-
-from .base import (
-    BaseModelArgs,
-    create_attention_mask,
-    create_ssm_mask,
-    scaled_dot_product_attention,
-)
-from .cache import ArraysCache, KVCache
-from .switch_layers import SwitchGLU
-
-
-@dataclass
-class ModelArgs(BaseModelArgs):
-    model_type: str
-    vocab_size: int
-    hidden_size: int
-    intermediate_size: int
-    moe_intermediate_size: int
-    num_hidden_layers: int
-    num_experts: int
-    num_experts_per_tok: int
-    norm_topk_prob: bool
-    num_attention_heads: int
-    num_key_value_heads: int
-    max_position_embeddings: int
-    use_expert_bias: bool
-    num_dense_layers: int
-    norm_eps: float
-    conv_bias: bool
-    conv_L_cache: int
-    rope_theta: float
-    full_attn_idxs: Optional[List[int]] = None
-    layer_types: Optional[List[str]] = None
-
-    def __post_init__(self):
-        if self.full_attn_idxs is None:
-            self.full_attn_idxs = [
-                i
-                for i, layer_type in enumerate(self.layer_types)
-                if layer_type == "full_attention"
-            ]
-
-
-class Attention(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-
-        dim = args.hidden_size
-        self.n_heads = n_heads = args.num_attention_heads
-        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
-
-        self.head_dim = head_dim = args.hidden_size // n_heads
-
-        self.scale = head_dim**-0.5
-
-        self.q_layernorm = nn.RMSNorm(head_dim, eps=args.norm_eps)
-        self.k_layernorm = nn.RMSNorm(head_dim, eps=args.norm_eps)
-
-        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False)
-        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
-        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
-        self.out_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
-
-        self.rope = nn.RoPE(
-            self.head_dim,
-            base=args.rope_theta,
-            traditional=False,
-        )
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-    ) -> mx.array:
-        B, L, D = x.shape
-
-        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
-
-        queries = self.q_layernorm(queries.reshape(B, L, self.n_heads, -1)).transpose(
-            0, 2, 1, 3
-        )
-        keys = self.k_layernorm(keys.reshape(B, L, self.n_kv_heads, -1)).transpose(
-            0, 2, 1, 3
-        )
-        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
-
-        if cache is not None:
-            queries = self.rope(queries, offset=cache.offset)
-            keys = self.rope(keys, offset=cache.offset)
-            keys, values = cache.update_and_fetch(keys, values)
-        else:
-            queries = self.rope(queries)
-            keys = self.rope(keys)
-
-        output = scaled_dot_product_attention(
-            queries, keys, values, cache=cache, mask=mask, scale=self.scale
-        )
-        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
-        return self.out_proj(output)
-
-
-class ShortConv(nn.Module):
-    def __init__(
-        self,
-        args: ModelArgs,
-        layer_idx: int,
-    ):
-        super().__init__()
-        self.args = args
-        self.layer_idx = layer_idx
-        self.L_cache = args.conv_L_cache
-        self.bias = args.conv_bias
-
-        self.conv = nn.Conv1d(
-            in_channels=args.hidden_size,
-            out_channels=args.hidden_size,
-            kernel_size=self.L_cache,
-            groups=args.hidden_size,
-            bias=self.bias,
-        )
-        self.in_proj = nn.Linear(args.hidden_size, 3 * args.hidden_size, bias=self.bias)
-        self.out_proj = nn.Linear(args.hidden_size, args.hidden_size, bias=self.bias)
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-    ):
-        BCx = self.in_proj(x)
-        B, C, x = mx.split(BCx, 3, axis=-1)
-        Bx = B * x
-        if mask is not None:
-            Bx = mx.where(mask[..., None], Bx, 0)
-        state = None
-        if cache is not None:
-            state = cache[0]
-        if state is None:
-            state = mx.zeros(
-                (Bx.shape[0], self.L_cache - 1, self.args.hidden_size), dtype=Bx.dtype
-            )
-
-        Bx = mx.concatenate([state, Bx], axis=-2)
-        if cache is not None:
-            cache[0] = Bx[:, -(self.L_cache - 1) :]
-        conv_out = self.conv(Bx)
-
-        y = C * conv_out
-        return self.out_proj(y)
-
-
-class MLP(nn.Module):
-    def __init__(self, config: ModelArgs, intermediate_size: Optional[int] = None):
-        super().__init__()
-        self.hidden_size = config.hidden_size
-        self.intermediate_size = (
-            config.intermediate_size if intermediate_size is None else intermediate_size
-        )
-        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
-        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
-
-    def __call__(self, x) -> mx.array:
-        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
-
-
-class Lfm2MoeSparseMoeBlock(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        dim = args.hidden_size
-        intermediate_size = args.moe_intermediate_size
-
-        self.num_experts = num_experts = args.num_experts
-        self.top_k = args.num_experts_per_tok
-        self.norm_topk_prob = args.norm_topk_prob
-        self.use_expert_bias = args.use_expert_bias
-
-        self.gate = nn.Linear(dim, num_experts, bias=False)
-        self.switch_mlp = SwitchGLU(dim, intermediate_size, num_experts)
-        if self.use_expert_bias:
-            self.expert_bias = mx.zeros((self.num_experts,))
-
-    def __call__(
-        self,
-        x: mx.array,
-    ):
-        gates = self.gate(x).astype(mx.float32)
-        gates = mx.softmax(gates, axis=-1)
-
-        if self.use_expert_bias:
-            gates += self.expert_bias
-
-        k = self.top_k
-        inds = mx.argpartition(gates, kth=-k, axis=-1)[..., -k:]
-
-        scores = mx.take_along_axis(gates, inds, axis=-1)
-        if self.norm_topk_prob:
-            scores /= mx.sum(scores, axis=-1, keepdims=True) + 1e-20
-        scores = scores.astype(x.dtype)
-
-        y = self.switch_mlp(x, inds)
-        y = (y * scores[..., None]).sum(axis=-2)
-
-        return y
-
-
-class Lfm2DecoderLayer(nn.Module):
-    def __init__(self, args: ModelArgs, layer_idx: int):
-        super().__init__()
-        self.is_attention_layer = layer_idx in args.full_attn_idxs
-
-        if self.is_attention_layer:
-            self.self_attn = Attention(args)
-        else:
-            self.conv = ShortConv(args, layer_idx)
-        self.feed_forward = (
-            MLP(
-                config=args,
-                intermediate_size=args.intermediate_size,
-            )
-            if layer_idx < args.num_dense_layers
-            else Lfm2MoeSparseMoeBlock(args)
-        )
-
-        self.operator_norm = nn.RMSNorm(args.hidden_size, eps=args.norm_eps)
-        self.ffn_norm = nn.RMSNorm(args.hidden_size, eps=args.norm_eps)
-
-    def __call__(
-        self,
-        x: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-    ) -> mx.array:
-
-        if self.is_attention_layer:
-            r = self.self_attn(self.operator_norm(x), mask=mask, cache=cache)
-        else:
-            r = self.conv(
-                self.operator_norm(x),
-                mask=mask,
-                cache=cache,
-            )
-        h = x + r
-        out = h + self.feed_forward(self.ffn_norm(h))
-        return out
-
-
-class Lfm2Model(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.vocab_size = args.vocab_size
-        self.num_hidden_layers = args.num_hidden_layers
-        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
-        self.layers = [
-            Lfm2DecoderLayer(args, layer_idx=i) for i in range(args.num_hidden_layers)
-        ]
-
-        self.embedding_norm = nn.RMSNorm(args.hidden_size, eps=args.norm_eps)
-
-        self.fa_idx = args.full_attn_idxs[0]
-        self.conv_idx = 0
-        for i in range(args.num_hidden_layers):
-            if i in args.full_attn_idxs:
-                self.conv_idx += 1
-            else:
-                break
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        cache=None,
-        input_embeddings: Optional[mx.array] = None,
-    ):
-        if input_embeddings is not None:
-            h = input_embeddings
-        else:
-            h = self.embed_tokens(inputs)
-
-        if cache is None:
-            cache = [None] * len(self.layers)
-
-        attn_mask = create_attention_mask(h, cache[self.fa_idx])
-        conv_mask = create_ssm_mask(h, cache[self.conv_idx])
-
-        for layer, c in zip(self.layers, cache):
-            mask = attn_mask if layer.is_attention_layer else conv_mask
-            h = layer(h, mask, cache=c)
-
-        return self.embedding_norm(h)
-
-
-class Model(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.model_type = args.model_type
-        self.model = Lfm2Model(args)
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        cache=None,
-        input_embeddings: Optional[mx.array] = None,
-    ):
-        out = self.model(inputs, cache, input_embeddings)
-        return self.model.embed_tokens.as_linear(out)
-
-    def sanitize(self, weights):
-        sanitized_weights = {}
-        for name, param in weights.items():
-            if "conv.weight" in name:
-                if param.shape[-1] > param.shape[1]:
-                    param = param.transpose(0, 2, 1)
-            replacements = {
-                "w1.weight": "gate_proj.weight",
-                "w2.weight": "down_proj.weight",
-                "w3.weight": "up_proj.weight",
-            }
-            for old, new in replacements.items():
-                if old in name:
-                    name = name.replace(old, new)
-            sanitized_weights[name] = param
-
-        for l in range(self.args.num_hidden_layers):
-            prefix = f"model.layers.{l}"
-            # Only sanitize MoE layer weights
-            for n in ["gate_proj", "down_proj", "up_proj"]:
-                if f"{prefix}.feed_forward.experts.0.{n}.weight" in sanitized_weights:
-                    to_join = [
-                        sanitized_weights.pop(
-                            f"{prefix}.feed_forward.experts.{e}.{n}.weight"
-                        )
-                        for e in range(self.args.num_experts)
-                    ]
-                    sanitized_weights[
-                        f"{prefix}.feed_forward.switch_mlp.{n}.weight"
-                    ] = mx.stack(to_join)
-        return sanitized_weights
-
-    @property
-    def layers(self):
-        return self.model.layers
-
-    def make_cache(self):
-        return [
-            KVCache() if l.is_attention_layer else ArraysCache(size=1)
-            for l in self.layers
-        ]
-
-    @property
-    def quant_predicate(self):
-        def predicate(path, _):
-            if path.endswith("feed_forward.gate"):
-                return {"group_size": 64, "bits": 8}
-            return True
-
-        return predicate
-
-    @property
-    def cast_predicate(self):
-        def predicate(k):
-            return "expert_bias" not in k
-
-        return predicate

mlx_lm/models/lille-130m.py

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

mlx_lm/models/llama.py

@@ -1,14 +1,10 @@
-# Copyright © 2023-2024 Apple Inc.
-
 from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Union
+from typing import 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 .cache import KVCache, RotatingKVCache
-from .rope_utils import initialize_rope
+from .base import BaseModelArgs, KVCache, create_additive_causal_mask
 
 
 @dataclass
@@ -20,8 +16,6 @@ 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
@@ -29,15 +23,18 @@ class ModelArgs(BaseModelArgs):
     rope_traditional: bool = False
     rope_scaling: Optional[Dict[str, Union[float, str]]] = None
     tie_word_embeddings: bool = True
-    layer_types: Optional[List[str]] = None
-    sliding_window: Optional[int] = None
 
     def __post_init__(self):
         if self.num_key_value_heads is None:
             self.num_key_value_heads = self.num_attention_heads
 
-        if self.layer_types is None:
-            self.layer_types = ["full_attention"] * self.num_hidden_layers
+        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):
@@ -48,8 +45,7 @@ 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
 
-        self.head_dim = head_dim = args.head_dim or args.hidden_size // n_heads
-
+        head_dim = args.hidden_size // n_heads
         self.scale = head_dim**-0.5
         if hasattr(args, "attention_bias"):
             attention_bias = args.attention_bias
@@ -61,19 +57,23 @@ 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)
 
-        self.rope = initialize_rope(
-            self.head_dim,
-            args.rope_theta,
-            args.rope_traditional,
-            args.rope_scaling,
-            args.max_position_embeddings,
+        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,
         )
 
     def __call__(
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
+        cache: Optional[KVCache] = None,
     ) -> mx.array:
         B, L, D = x.shape
 
@@ -92,10 +92,9 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        output = scaled_dot_product_attention(
-            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        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.o_proj(output)
 
@@ -120,11 +119,10 @@ class MLP(nn.Module):
 
 
 class TransformerBlock(nn.Module):
-    def __init__(self, args: ModelArgs, use_sliding: bool = False):
+    def __init__(self, args: ModelArgs):
         super().__init__()
         self.num_attention_heads = args.num_attention_heads
         self.hidden_size = args.hidden_size
-        self.use_sliding = use_sliding
         self.self_attn = Attention(args)
         self.mlp = MLP(args)
         self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
@@ -137,7 +135,7 @@ class TransformerBlock(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
+        cache: Optional[KVCache] = None,
     ) -> mx.array:
         r = self.self_attn(self.input_layernorm(x), mask, cache)
         h = x + r
@@ -152,45 +150,32 @@ class LlamaModel(nn.Module):
         self.args = args
         self.vocab_size = args.vocab_size
         self.num_hidden_layers = args.num_hidden_layers
-        self.layer_types = args.layer_types
-        self.sliding_window = args.sliding_window
         assert self.vocab_size > 0
         self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
         self.layers = [
-            TransformerBlock(args=args, use_sliding=layer_type == "sliding_attention")
-            for layer_type in self.layer_types
+            TransformerBlock(args=args) for _ in range(args.num_hidden_layers)
         ]
         self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
-        self.fa_idx = self.layer_types.index("full_attention")
-        self.swa_idx = None
-        for e, l in enumerate(self.layers):
-            if l.use_sliding:
-                self.swa_idx = e
-                break
 
     def __call__(
         self,
         inputs: mx.array,
         cache=None,
-        input_embeddings: Optional[mx.array] = None,
     ):
-        if input_embeddings is not None:
-            h = input_embeddings
-        else:
-            h = self.embed_tokens(inputs)
+        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 cache is None:
             cache = [None] * len(self.layers)
 
-        fa_mask = create_attention_mask(h, cache[self.fa_idx])
-        if self.swa_idx is not None:
-            swa_mask = create_attention_mask(
-                h, cache[self.swa_idx], window_size=self.sliding_window
-            )
-
-        for layer, cache in zip(self.layers, cache):
-            mask = swa_mask if layer.use_sliding else fa_mask
-            h = layer(h, mask, cache=cache)
+        for layer, c in zip(self.layers, cache):
+            h = layer(h, mask, cache=c)
 
         return self.norm(h)
 
@@ -208,9 +193,8 @@ class Model(nn.Module):
         self,
         inputs: mx.array,
         cache=None,
-        input_embeddings: Optional[mx.array] = None,
     ):
-        out = self.model(inputs, cache, input_embeddings)
+        out = self.model(inputs, cache)
         if self.args.tie_word_embeddings:
             out = self.model.embed_tokens.as_linear(out)
         else:
@@ -219,23 +203,18 @@ class Model(nn.Module):
 
     def sanitize(self, weights):
         # Remove unused precomputed rotary freqs
-        weights = {
+        return {
             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
 
-    def make_cache(self):
-        return [
-            (
-                RotatingKVCache(max_size=self.model.sliding_window)
-                if layer.use_sliding
-                else KVCache()
-            )
-            for layer in self.layers
-        ]
+    @property
+    def 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

@@ -1,324 +0,0 @@
-# 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_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
-        assert self.top_k == 1, "Only 1 expert per token supported"
-        self.num_experts = args.num_local_experts
-        self.experts = SwitchGLU(
-            args.hidden_size, args.intermediate_size, self.num_experts
-        )
-        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,
-        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 cache is None:
-            cache = [None] * len(self.layers)
-
-        global_mask = create_attention_mask(h, cache[3])
-
-        for idx, (layer, c) in enumerate(zip(self.layers, cache)):
-            use_chunked_attention = (idx + 1) % 4 != 0
-            mask = chunk_mask if use_chunked_attention else global_mask
-            h = layer(h, 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,
-        cache=None,
-    ):
-        out = self.model(inputs, 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,
-        cache=None,
-    ):
-        return self.language_model(inputs, 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/llama4_text.py

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

mlx_lm/models/longcat_flash.py

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

mlx_lm/models/mamba.py

@@ -1,220 +0,0 @@
-# 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 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 = nn.Conv1d(
-            in_channels=self.intermediate_size,
-            out_channels=self.intermediate_size,
-            kernel_size=self.conv_kernel_size,
-            groups=self.intermediate_size,
-            bias=self.use_conv_bias,
-            padding=0,
-        )
-
-        self.x_proj = nn.Linear(
-            self.intermediate_size,
-            self.time_step_rank + 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)
-        K = self.conv_kernel_size
-        if conv_cache is not None:
-            x_full = mx.concatenate([conv_cache, x], axis=1)
-        else:
-            x_full = mx.pad(x, [(0, 0), (K - 1, 0), (0, 0)])
-        conv_out = self.conv1d(x_full)
-        new_conv_cache = x_full[:, -(K - 1) :, :]
-        x = nn.silu(conv_out)
-        A = -mx.exp(self.A_log)
-        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 make_cache(self):
-        return [MambaCache() for _ in range(len(self.layers))]
-
-    @property
-    def layers(self):
-        return self.backbone.layers
-
-    def sanitize(self, weights):
-        for k, v in weights.items():
-            if "conv1d.weight" in k and v.shape[-1] != 1:
-                weights[k] = v.moveaxis(2, 1)
-        return weights

mlx_lm/models/mamba2.py

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

mlx_lm/models/mimo.py

@@ -1,193 +0,0 @@
-# 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,
-        cache=None,
-    ):
-        h = self.embed_tokens(inputs)
-
-        if cache is None:
-            cache = [None] * len(self.layers)
-
-        mask = create_attention_mask(h, cache[0])
-
-        for layer, c in zip(self.layers, cache):
-            h = layer(h, mask, c)
-
-        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,
-        cache=None,
-    ):
-        out = self.model(inputs, cache)
-
-        if self.args.tie_word_embeddings:
-            out = self.model.embed_tokens.as_linear(out)
-        else:
-            out = self.lm_head(out)
-
-        return out
-
-    def sanitize(self, weights):
-        if self.args.tie_word_embeddings:
-            weights.pop("lm_head.weight", None)
-
-        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,13 +1,11 @@
-# Copyright © 2023-2025 Apple Inc.
-
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Union
+from typing import Dict, Optional, Tuple, Union
 
 import mlx.core as mx
 import mlx.nn as nn
+import numpy as np
 
-from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
-from .rope_utils import initialize_rope
+from .base import BaseModelArgs
 
 
 @dataclass
@@ -23,7 +21,6 @@ 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
@@ -69,19 +66,24 @@ class Attention(nn.Module):
             self.num_heads * self.head_dim, self.hidden_size, bias=False
         )
 
-        self.rope = initialize_rope(
-            self.head_dim,
-            args.rope_theta,
-            args.rope_traditional,
-            args.rope_scaling,
-            args.max_position_embeddings,
+        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,
         )
 
     def __call__(
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = None,
     ):
         B, L, _ = x.shape
 
@@ -101,8 +103,8 @@ class Attention(nn.Module):
             queries = self.rope(queries)
             keys = self.rope(keys)
 
-        attn_output = scaled_dot_product_attention(
-            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        attn_output = mx.fast.scaled_dot_product_attention(
+            queries, keys, values, scale=self.scale, mask=mask
         )
 
         attn_output = attn_output.transpose(0, 2, 1, 3).reshape(B, L, -1)
@@ -131,12 +133,12 @@ class DecoderLayer(nn.Module):
         self,
         x: mx.array,
         mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
+        cache: Optional[Tuple[mx.array, mx.array]] = 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)
+        h = x + r * (self.scale_depth / np.sqrt(self.num_hidden_layers))
         r = self.mlp(self.post_attention_layernorm(h))
-        out = h + r * (self.scale_depth / self.num_hidden_layers**0.5)
+        out = h + r * (self.scale_depth / np.sqrt(self.num_hidden_layers))
         return out
 
 
@@ -158,11 +160,14 @@ class MiniCPMModel(nn.Module):
     ):
         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 cache is None:
             cache = [None] * len(self.layers)
 
-        mask = create_attention_mask(h, cache[0])
-
         for layer, c in zip(self.layers, cache):
             h = layer(h, mask, c)
 
@@ -201,3 +206,11 @@ 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

@@ -1,250 +0,0 @@
-# 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/minimax.py

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

mlx_lm/models/ministral3.py

@@ -1,264 +0,0 @@
-# Copyright © 2023-2024 Apple Inc.
-
-from dataclasses import dataclass
-from typing import Any, Dict, List, Optional, Union
-
-import mlx.core as mx
-import mlx.nn as nn
-
-from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
-from .cache import KVCache, RotatingKVCache
-from .rope_utils import initialize_rope
-
-
-@dataclass
-class ModelArgs(BaseModelArgs):
-    model_type: str
-    hidden_size: int
-    num_hidden_layers: int
-    intermediate_size: int
-    num_attention_heads: int
-    rms_norm_eps: float
-    vocab_size: int
-    head_dim: Optional[int] = None
-    max_position_embeddings: Optional[int] = None
-    num_key_value_heads: Optional[int] = None
-    rope_parameters: Optional[Dict[str, Union[float, str]]] = None
-    tie_word_embeddings: bool = True
-    layer_types: Optional[List[str]] = None
-    sliding_window: Optional[int] = None
-
-    def __post_init__(self):
-        if self.num_key_value_heads is None:
-            self.num_key_value_heads = self.num_attention_heads
-
-        if self.layer_types is None:
-            self.layer_types = ["full_attention"] * self.num_hidden_layers
-
-
-def _get_llama_4_attn_scale(
-    start: int, stop: int, beta: float, max_position_embeddings: int
-):
-    scaling = 1 + beta * mx.log(
-        1 + mx.floor(mx.arange(start, stop) / max_position_embeddings)
-    )
-    return scaling[:, None]
-
-
-class Attention(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-
-        dim = args.hidden_size
-        self.n_heads = n_heads = args.num_attention_heads
-        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
-
-        self.head_dim = head_dim = args.head_dim or args.hidden_size // n_heads
-
-        self.scale = head_dim**-0.5
-
-        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=False)
-        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
-        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=False)
-        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=False)
-
-        self.rope = initialize_rope(
-            self.head_dim,
-            args.rope_parameters["rope_theta"],
-            False,
-            args.rope_parameters,
-            args.max_position_embeddings,
-        )
-
-    def __call__(
-        self,
-        x: mx.array,
-        attn_scale: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-    ) -> mx.array:
-        B, L, D = x.shape
-
-        queries, keys, values = self.q_proj(x), self.k_proj(x), self.v_proj(x)
-
-        # Prepare the queries, keys and values for the attention computation
-        queries = queries.reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
-        keys = keys.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
-        values = values.reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
-
-        offset = 0
-        if cache is not None:
-            offset = cache.offset
-            queries = self.rope(queries, offset=offset)
-            keys = self.rope(keys, offset=offset)
-            keys, values = cache.update_and_fetch(keys, values)
-        else:
-            queries = self.rope(queries)
-            keys = self.rope(keys)
-        queries = queries * attn_scale
-        output = scaled_dot_product_attention(
-            queries, keys, values, cache=cache, scale=self.scale, mask=mask
-        )
-
-        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
-        return self.o_proj(output)
-
-
-class MLP(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-
-        dim = args.hidden_size
-        hidden_dim = args.intermediate_size
-        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
-        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)
-        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
-
-    def __call__(self, x) -> mx.array:
-        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
-
-
-class TransformerBlock(nn.Module):
-    def __init__(self, args: ModelArgs, use_sliding: bool = False):
-        super().__init__()
-        self.num_attention_heads = args.num_attention_heads
-        self.hidden_size = args.hidden_size
-        self.use_sliding = use_sliding
-        self.self_attn = Attention(args)
-        self.mlp = MLP(args)
-        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
-        self.post_attention_layernorm = nn.RMSNorm(
-            args.hidden_size, eps=args.rms_norm_eps
-        )
-        self.args = args
-
-    def __call__(
-        self,
-        x: mx.array,
-        attn_scale: mx.array,
-        mask: Optional[mx.array] = None,
-        cache: Optional[Any] = None,
-    ) -> mx.array:
-        r = self.self_attn(self.input_layernorm(x), attn_scale, mask, cache)
-        h = x + r
-        r = self.mlp(self.post_attention_layernorm(h))
-        out = h + r
-        return out
-
-
-class LanguageModel(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.vocab_size = args.vocab_size
-        self.num_hidden_layers = args.num_hidden_layers
-        self.layer_types = args.layer_types
-        self.sliding_window = args.sliding_window
-        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
-        self.layers = [
-            TransformerBlock(args=args, use_sliding=layer_type == "sliding_attention")
-            for layer_type in self.layer_types
-        ]
-        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
-        self.fa_idx = self.layer_types.index("full_attention")
-        self.swa_idx = None
-        for e, l in enumerate(self.layers):
-            if l.use_sliding:
-                self.swa_idx = e
-                break
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        cache=None,
-        input_embeddings: Optional[mx.array] = None,
-    ):
-        if input_embeddings is not None:
-            h = input_embeddings
-        else:
-            h = self.embed_tokens(inputs)
-
-        if cache is None:
-            cache = [None] * len(self.layers)
-            offset = 0
-        else:
-            offset = cache[0].offset
-
-        fa_mask = create_attention_mask(h, cache[self.fa_idx])
-        if self.swa_idx is not None:
-            swa_mask = create_attention_mask(
-                h, cache[self.swa_idx], window_size=self.sliding_window
-            )
-
-        attn_scale = _get_llama_4_attn_scale(
-            offset,
-            offset + inputs.shape[1],
-            self.args.rope_parameters["llama_4_scaling_beta"],
-            self.args.rope_parameters["original_max_position_embeddings"],
-        ).astype(h.dtype)
-
-        for layer, cache in zip(self.layers, cache):
-            mask = swa_mask if layer.use_sliding else fa_mask
-            h = layer(h, attn_scale, mask, cache=cache)
-
-        return self.norm(h)
-
-
-class Model(nn.Module):
-    def __init__(self, args: ModelArgs):
-        super().__init__()
-        self.args = args
-        self.model_type = args.model_type
-        self.model = LanguageModel(args)
-        if not args.tie_word_embeddings:
-            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        cache=None,
-        input_embeddings: Optional[mx.array] = None,
-    ):
-        out = self.model(inputs, cache, input_embeddings)
-        if self.args.tie_word_embeddings:
-            out = self.model.embed_tokens.as_linear(out)
-        else:
-            out = self.lm_head(out)
-        return out
-
-    def sanitize(self, weights):
-        # Remove unused precomputed rotary freqs
-        weights = {
-            k: v for k, v in weights.items() if "self_attn.rotary_emb.inv_freq" not in k
-        }
-        if self.args.tie_word_embeddings:
-            weights.pop("lm_head.weight", None)
-
-        new_weights = {}
-        for k, v in weights.items():
-            if "weight_scale_inv" in k:
-                scale_inv = v
-                wk = k.replace("_scale_inv", "")
-                weight = weights[wk]
-                new_weights[wk] = weight * scale_inv
-            elif "activation_scale" in k:
-                continue
-            elif k not in new_weights:
-                new_weights[k] = v
-        weights = new_weights
-
-        return weights
-
-    @property
-    def layers(self):
-        return self.model.layers
-
-    def make_cache(self):
-        return [
-            (
-                RotatingKVCache(max_size=self.model.sliding_window)
-                if layer.use_sliding
-                else KVCache()
-            )
-            for layer in self.layers
-        ]

mlx_lm/models/mistral3.py

@@ -1,58 +0,0 @@
-# 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, ministral3
-from .base import BaseModelArgs
-
-
-@dataclass
-class ModelArgs(BaseModelArgs):
-    model_type: str
-    text_config: dict
-
-    def __post_init__(self):
-        if "tie_word_embeddings" not in self.text_config:
-            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
-        if args.text_config.get("model_type") == "ministral3":
-            self.language_model = ministral3.Model(
-                ministral3.ModelArgs.from_dict(args.text_config)
-            )
-        else:
-            self.language_model = llama.Model(
-                llama.ModelArgs.from_dict(args.text_config)
-            )
-
-    def __call__(
-        self,
-        inputs: mx.array,
-        cache=None,
-        input_embeddings: Optional[mx.array] = None,
-    ):
-        return self.language_model(
-            inputs, cache=cache, input_embeddings=input_embeddings
-        )
-
-    def sanitize(self, weights):
-        weights = tree_unflatten(list(weights.items()))
-        weights.pop("vision_tower", None)
-        weights.pop("multi_modal_projector", None)
-        lm_weights = dict(tree_flatten(weights["language_model"]))
-        weights["language_model"] = self.language_model.sanitize(lm_weights)
-        return dict(tree_flatten(weights))
-
-    @property
-    def layers(self):
-        return self.language_model.model.layers