Add glm4 moe lite model (#776)

* Add glm4 moe lite model

* Format glm4_moe_lite with black

* version

---------

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

.circleci/config.yml

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

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

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

.github/workflows/pull_request.yml

@@ -0,0 +1,43 @@
+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: |
+          curl -o test_data.zip -L https://github.com/ml-explore/mlx-lm/releases/download/test_data/test_data.zip
+          unzip test_data.zip
+          HF_HOME="." python -m xmlrunner discover -v tests -o test-results/

.github/workflows/release.yml

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

ACKNOWLEDGMENTS.md

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

README.md

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

benchmarks/server_benchmark.py

@@ -0,0 +1,348 @@
+"""
+Spin up the local server:
+
+    mlx_lm.server
+
+Then run the benchmark:
+
+    python server_benchmark.py --concurrency 4
+"""
+
+import argparse
+import asyncio
+import json
+import math
+import time
+from collections import defaultdict
+from itertools import cycle
+from typing import Any, Dict, List, Optional, Tuple
+
+import aiohttp
+from tqdm import tqdm
+
+# Default prompts if no file is provided
+DEFAULT_PROMPTS = [
+    "Explain quantum computing in simple terms.",
+    "What are the main differences between Python and JavaScript?",
+    "Describe the process of photosynthesis in plants.",
+    "How does a neural network learn from data?",
+    "What is the significance of the Turing test in AI?",
+    "Explain the concept of blockchain technology.",
+    "What causes seasons on Earth?",
+    "How do vaccines work in the human body?",
+    "Describe the water cycle and its importance.",
+    "What is the theory of relativity proposed by Einstein?",
+    "How do electric cars help reduce carbon emissions?",
+    "What are the key features of a market economy?",
+    "Explain how DNA replication works in cells.",
+    "What is machine learning and its real-world applications?",
+    "Describe the structure and function of the human heart.",
+]
+
+
+def tokens_per_second(tokens):
+    start = math.floor(tokens[0])
+    stop = math.ceil(tokens[-1])
+    n_bins = int(stop - start) * 10
+    bins = [0] * n_bins
+    for t in tokens:
+        bins[int(n_bins * (t - start) / (stop - start))] += 1
+
+    result = []
+
+    ms = 0
+    cnt = 0
+    for i, b in enumerate(bins):
+        ms += b
+        if cnt == 10:
+            ms -= bins[i - 10]
+        else:
+            cnt += 1
+
+        result.append(10 * ms / cnt)
+
+    times = [start]
+    while times[-1] < stop:
+        times.append(times[-1] + 0.1)
+
+    return times, result
+
+
+def plot_generation(times, tokens_per_sec, start=None, interval=1.0, width=50):
+    c = "█"
+    start = start or times[0]
+    stop = times[-1]
+
+    bar_times = [start]
+    while bar_times[-1] < stop:
+        bar_times.append(bar_times[-1] + interval)
+
+    bar_values = [[] for _ in bar_times]
+    bar_idx = 0
+
+    for t, v in zip(times, tokens_per_sec):
+        while t > bar_times[bar_idx] + interval:
+            bar_idx += 1
+        bar_values[bar_idx].append(v)
+
+    bar_values = [sum(v) / len(v) if v else 0 for v in bar_values]
+    m = max(bar_values)
+
+    for t, v in zip(bar_times, bar_values):
+        t = t - start
+        b = c * int(v * width / m)
+        print(f"{t:3.2f} {b} ({v})")
+
+
+def percentile(data, percent):
+    if not data:
+        return 0
+    data = sorted(data)
+    k = (len(data) - 1) * percent / 100
+    f = math.floor(k)
+    c = math.ceil(k)
+    return (
+        data[int(f)]
+        if f == c
+        else data[int(f)] + (data[int(c)] - data[int(f)]) * (k - f)
+    )
+
+
+def median(data):
+    return percentile(data, 50)
+
+
+async def make_request(
+    session: aiohttp.ClientSession,
+    url: str,
+    api_key: str,
+    model: str,
+    prompt: str,
+    max_tokens: int,
+) -> Tuple[bool, float, list]:
+    """
+    Make a single streaming API request and return
+
+        - whether the request succeeded
+        - the request start time
+        - the time of every generated token
+    """
+    payload = {
+        "model": model,
+        "messages": [{"role": "user", "content": prompt}],
+        "max_tokens": max_tokens,
+        "stream": True,
+    }
+    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
+
+    start_time = time.perf_counter()
+    tokens = []
+
+    try:
+        async with session.post(url, json=payload, headers=headers) as response:
+            if response.status != 200:
+                error_body = await response.text()
+                print(f"Error {response.status}: {error_body}")
+                return (False, 0, [])
+
+            # Process streaming response
+            async for chunk in response.content:
+                if chunk:
+                    chunk_str = chunk.decode("utf-8").strip()
+                    if chunk_str.startswith("data:"):
+                        data_str = chunk_str[5:].strip()
+                        if data_str == "[DONE]":
+                            break
+
+                        try:
+                            data = json.loads(data_str)
+                            if choices := data.get("choices", False):
+                                if choices[0].get("finish_reason") != "length":
+                                    tokens.append(time.perf_counter())
+                        except json.JSONDecodeError:
+                            continue
+
+            return (bool(tokens), start_time, tokens)
+
+    except Exception as e:
+        print(f"Request failed: {str(e)}")
+        return (False, 0, [])
+
+
+async def run_benchmark(
+    url: str,
+    api_key: str,
+    model: str,
+    max_tokens: int,
+    concurrency: int,
+    total_requests: int,
+    prompts: List[str],
+) -> Dict[str, Any]:
+    prompt_cycle = cycle(prompts)
+    semaphore = asyncio.Semaphore(concurrency)
+    results = []
+    request_times = []
+    bar = tqdm(total=total_requests)
+
+    async def worker():
+        async with semaphore:
+            prompt = next(prompt_cycle)
+            result = await make_request(
+                session, url, api_key, model, prompt, max_tokens
+            )
+            bar.update(1)
+            return result
+
+    async with aiohttp.ClientSession() as session:
+        tasks = []
+        for _ in range(total_requests):
+            task = asyncio.create_task(worker())
+            tasks.append(task)
+            await asyncio.sleep(0.01)  # Stagger requests slightly
+
+        for task in tasks:
+            result = await task
+            results.append(result)
+        bar.close()
+
+    successful_requests = [r for r in results if r[0]]
+    total_tokens = sum(len(r[2]) for r in successful_requests)
+
+    # Gather all the tokens generated with their corresponding timestamps
+    all_tokens = []
+    for r in successful_requests:
+        all_tokens.extend(r[2])
+    all_tokens.sort()
+    full_generation = tokens_per_second(all_tokens)
+    start = min(r[1] for r in successful_requests)
+
+    # Aggregate metrics
+    metrics = {
+        "total_requests": total_requests,
+        "successful_requests": len(successful_requests),
+        "failed_requests": total_requests - len(successful_requests),
+        "total_tokens": total_tokens,
+        "total_time": all_tokens[-1] - start,
+        "aggregate_tokens_per_sec": median(full_generation[1]),
+        "per_request": [],
+        "start": start,
+        "full_generation": full_generation,
+    }
+
+    # Per-request metrics
+    for i, (_, start, tokens) in enumerate(successful_requests):
+        metrics["per_request"].append(
+            {
+                "request_id": i + 1,
+                "time_to_first_token": tokens[0] - start,
+                "total_time": tokens[-1] - start,
+                "tokens_received": len(tokens),
+                "tokens_per_sec": median(tokens_per_second(tokens)[1]),
+            }
+        )
+
+    # Calculate percentiles
+    ttft_values = [m["time_to_first_token"] for m in metrics["per_request"]]
+    tps_values = [m["tokens_per_sec"] for m in metrics["per_request"]]
+
+    metrics["aggregate_metrics"] = {
+        "time_to_first_token": {
+            "min": min(ttft_values) if ttft_values else 0,
+            "max": max(ttft_values) if ttft_values else 0,
+            "avg": sum(ttft_values) / len(ttft_values) if ttft_values else 0,
+            "p95": percentile(ttft_values, 95) if ttft_values else 0,
+        },
+        "tokens_per_sec": {
+            "min": min(tps_values) if tps_values else 0,
+            "max": max(tps_values) if tps_values else 0,
+            "avg": sum(tps_values) / len(tps_values) if tps_values else 0,
+            "p95": percentile(tps_values, 95) if tps_values else 0,
+        },
+    }
+
+    return metrics
+
+
+def main():
+    parser = argparse.ArgumentParser(description="LLM API Benchmark Tool")
+    parser.add_argument(
+        "--url",
+        default="http://localhost:8080/v1/chat/completions",
+        help="Chat completions API endpoint URL",
+    )
+    parser.add_argument("--api-key", default="none", help="API key")
+    parser.add_argument("--model", default="default_model", help="Model name")
+    parser.add_argument(
+        "--max-tokens", type=int, default=100, help="Max tokens to generate"
+    )
+    parser.add_argument(
+        "--concurrency", type=int, default=1, help="Number of concurrent requests"
+    )
+    parser.add_argument(
+        "--total-requests", type=int, default=10, help="Total requests to make"
+    )
+    parser.add_argument("--prompt-file", help="File containing prompts (one per line)")
+    parser.add_argument("--output", help="Output file for results (JSON format)")
+
+    args = parser.parse_args()
+
+    # Load prompts
+    if args.prompt_file:
+        with open(args.prompt_file, "r") as f:
+            prompts = [line.strip() for line in f if line.strip()]
+    else:
+        prompts = DEFAULT_PROMPTS
+
+    print(
+        f"Starting benchmark with {args.concurrency} concurrency and {args.total_requests} total requests..."
+    )
+    start_time = time.perf_counter()
+
+    # Run benchmark
+    results = asyncio.run(
+        run_benchmark(
+            url=args.url,
+            api_key=args.api_key,
+            model=args.model,
+            max_tokens=args.max_tokens,
+            concurrency=args.concurrency,
+            total_requests=args.total_requests,
+            prompts=prompts,
+        )
+    )
+
+    duration = time.perf_counter() - start_time
+    print(f"\nBenchmark completed in {duration:.2f} seconds")
+    print(
+        f"Successful requests: {results['successful_requests']}/{args.total_requests}"
+    )
+    print(f"Total tokens generated: {results['total_tokens']}")
+    print(f"Aggregate tokens/sec: {results['aggregate_tokens_per_sec']:.2f}")
+
+    # Print summary
+    if results["successful_requests"] > 0:
+        ttft = results["aggregate_metrics"]["time_to_first_token"]
+        tps = results["aggregate_metrics"]["tokens_per_sec"]
+
+        print("\nTime to First Token (seconds):")
+        print(
+            f"  Min: {ttft['min']:.4f} | Max: {ttft['max']:.4f} | Avg: {ttft['avg']:.4f} | P95: {ttft['p95']:.4f}"
+        )
+
+        print("\nTokens per Second (per request):")
+        print(
+            f"  Min: {tps['min']:.2f} | Max: {tps['max']:.2f} | Avg: {tps['avg']:.2f} | P95: {tps['p95']:.2f}"
+        )
+
+        print()
+        plot_generation(*results["full_generation"], results["start"])
+
+    # Save results
+    if args.output:
+        with open(args.output, "w") as f:
+            json.dump(results, f, indent=2)
+        print(f"\nResults saved to {args.output}")
+
+
+if __name__ == "__main__":
+    main()

mlx_lm/BENCHMARKS.md

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

mlx_lm/LEARNED_QUANTS.md

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

mlx_lm/LORA.md

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

mlx_lm/__init__.py

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

mlx_lm/__main__.py

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

mlx_lm/_version.py

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

mlx_lm/benchmark.py

@@ -4,13 +4,9 @@ import argparse
 
 import mlx.core as mx
 
-from mlx_lm import batch_generate, stream_generate
+from mlx_lm import batch_generate, load, stream_generate
 from mlx_lm.generate import DEFAULT_MODEL
-from mlx_lm.tokenizer_utils import load_tokenizer
-from mlx_lm.utils import (
-    fetch_from_hub,
-    get_model_path,
-)
+from mlx_lm.utils import pipeline_load, sharded_load
 
 
 def setup_arg_parser():
@@ -25,11 +21,6 @@ def setup_arg_parser():
         ),
         default=None,
     )
-    parser.add_argument(
-        "--trust-remote-code",
-        action="store_true",
-        help="Enable trusting remote code for tokenizer",
-    )
     parser.add_argument(
         "--prompt-tokens",
         "-p",
@@ -58,6 +49,11 @@ def setup_arg_parser():
         help="Number of timing trials",
         type=int,
     )
+    parser.add_argument(
+        "--pipeline",
+        action="store_true",
+        help="Use pipelining instead of tensor parallelism",
+    )
     return parser
 
 
@@ -66,22 +62,34 @@ def main():
     args = parser.parse_args()
     mx.random.seed(0)
 
+    group = mx.distributed.init()
+    rank = group.rank()
+    pipeline_group = group if args.pipeline else None
+    tensor_group = group if not args.pipeline else None
+
+    def rprint(*args, **kwargs):
+        if rank == 0:
+            print(*args, **kwargs)
+
     model_path = args.model or DEFAULT_MODEL
 
-    model_path, _ = get_model_path(model_path, revision=None)
-    model, config, _ = fetch_from_hub(model_path, trust_remote_code=True)
-    tokenizer = load_tokenizer(
-        model_path,
-        eos_token_ids=[],  # Empty to avoid early stopping
-        tokenizer_config_extra={"trust_remote_code": True},
-    )
+    if group.size() > 1:
+        model, tokenizer, config = sharded_load(
+            model_path, pipeline_group, tensor_group, return_config=True
+        )
+    else:
+        model, tokenizer, config = load(
+            model_path, return_config=True, tokenizer_config={"trust_remote_code": True}
+        )
+
+    # Empty to avoid early stopping
+    tokenizer._eos_token_ids = {}
 
     prompt_tokens = args.prompt_tokens
     generation_tokens = args.generation_tokens
     batch_size = args.batch_size
-    prompts = mx.random.randint(
-        0, config["vocab_size"], (batch_size, prompt_tokens)
-    ).tolist()
+    vocab_size = config.get("vocab_size") or config["text_config"]["vocab_size"]
+    prompts = mx.random.randint(0, vocab_size, (batch_size, prompt_tokens)).tolist()
     prompt = prompts[0]
 
     def single_bench():
@@ -96,20 +104,23 @@ def main():
             model, tokenizer, prompts, max_tokens=generation_tokens
         ).stats
 
-    _bench = batch_bench
+    if batch_size == 1:
+        _bench = single_bench
+    else:
+        _bench = batch_bench
 
-    print("Running warmup..")
+    rprint("Running warmup..")
     _bench()
 
     report_keys = ["prompt_tps", "generation_tps", "peak_memory"]
-    print(f"Timing with {prompt_tokens=}, {generation_tokens=}, {batch_size=}.")
+    rprint(f"Timing with {prompt_tokens=}, {generation_tokens=}, {batch_size=}.")
     responses = []
     for i in range(args.num_trials):
         response = _bench()
         responses.append(response)
         results = [(k, getattr(response, k)) for k in report_keys]
         results = [f"{k}={v:.3f}" for k, v in results]
-        print(f"Trial {i+1}:  " + ", ".join(results))
+        rprint(f"Trial {i+1}:  " + ", ".join(results))
 
     def avg(k):
         vals = (getattr(response, k) for response in responses)
@@ -117,7 +128,7 @@ def main():
 
     results = [(k, avg(k)) for k in report_keys]
     results = [f"{k}={v:.3f}" for k, v in results]
-    print(f"Averages: " + ", ".join(results))
+    rprint(f"Averages: " + ", ".join(results))
 
 
 if __name__ == "__main__":

mlx_lm/cache_prompt.py

@@ -41,16 +41,6 @@ def setup_arg_parser():
         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,
@@ -107,14 +97,12 @@ def main():
 
     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:
+    if tokenizer.has_chat_template:
         messages = [{"role": "user", "content": args.prompt}]
         prompt = tokenizer.apply_chat_template(
-            messages, add_generation_prompt=False, continue_final_message=True
+            messages,
+            add_generation_prompt=False,
+            continue_final_message=True,
         )
 
     else:
@@ -153,7 +141,6 @@ def main():
     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)
 

mlx_lm/chat.py

@@ -7,7 +7,7 @@ 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
+from .utils import load, sharded_load
 
 DEFAULT_TEMP = 0.0
 DEFAULT_TOP_P = 1.0
@@ -79,6 +79,11 @@ def setup_arg_parser():
         default=None,
         help="System prompt to be used for the chat template",
     )
+    parser.add_argument(
+        "--pipeline",
+        action="store_true",
+        help="Use pipelining instead of tensor parallelism",
+    )
     return parser
 
 
@@ -86,28 +91,42 @@ def main():
     parser = setup_arg_parser()
     args = parser.parse_args()
 
+    group = mx.distributed.init()
+    rank = group.rank()
+    pipeline_group = group if args.pipeline else None
+    tensor_group = group if not args.pipeline else None
+
+    def rprint(*args, **kwargs):
+        if rank == 0:
+            print(*args, **kwargs)
+
     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
-        },
-    )
+    if group.size() > 1:
+        if args.adapter_path:
+            parser.error("Adapters not supported in distributed mode")
+        model, tokenizer = sharded_load(args.model, pipeline_group, tensor_group)
+    else:
+        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")
+        rprint("The command list:")
+        rprint("- 'q' to exit")
+        rprint("- 'r' to reset the chat")
+        rprint("- 'h' to display these commands")
 
-    print(f"[INFO] Starting chat session with {args.model}.")
+    rprint(f"[INFO] Starting chat session with {args.model}.")
     print_help()
     prompt_cache = make_prompt_cache(model, args.max_kv_size)
     while True:
-        query = input(">> ")
+        query = input(">> " if rank == 0 else "")
         if query == "q":
             break
         if query == "r":
@@ -120,7 +139,10 @@ def main():
         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)
+        prompt = tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+        )
         for response in stream_generate(
             model,
             tokenizer,
@@ -137,8 +159,8 @@ def main():
             ),
             prompt_cache=prompt_cache,
         ):
-            print(response.text, flush=True, end="")
-        print()
+            rprint(response.text, flush=True, end="")
+        rprint()
 
 
 if __name__ == "__main__":

mlx_lm/chat_templates/deepseek_v32.py

@@ -0,0 +1,332 @@
+# Copyright © 2025 Apple Inc.
+
+import copy
+import json
+import re
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+TOOLS_SYSTEM_TEMPLATE = """## Tools
+
+You have access to a set of tools you can use to answer the user's question.
+You can invoke functions by writing a "<{dsml_token}function_calls>" block like the following as part of your reply to the user:
+<{dsml_token}function_calls>
+<{dsml_token}invoke name="$FUNCTION_NAME">
+<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE</{dsml_token}parameter>
+...
+</{dsml_token}invoke>
+<{dsml_token}invoke name="$FUNCTION_NAME2">
+...
+</{dsml_token}invoke>
+</{dsml_token}function_calls>
+
+String and scalar parameters should be specified as is without any escaping or quotes, while lists and objects should use JSON format. The "string" attribute should be set to "true" for string type parameters and "false" for other types (numbers, booleans, arrays, objects).
+
+If the thinking_mode is enabled, then after function results you should strongly consider outputting a thinking block. Here is an example:
+
+<{dsml_token}function_calls>
+...
+</{dsml_token}function_calls>
+
+<function_results>
+...
+</function_results>
+
+{thinking_start_token}...thinking about results{thinking_end_token}
+
+Here are the functions available in JSONSchema format:
+<functions>
+{tool_schemas}
+</functions>
+"""
+
+bos_token: str = "<｜begin▁of▁sentence｜>"
+eos_token: str = "<｜end▁of▁sentence｜>"
+thinking_start_token: str = "<think>"
+thinking_end_token: str = "</think>"
+dsml_token: str = "｜DSML｜"
+system_msg_template: str = "{content}"
+user_msg_template: str = "<｜User｜>{content}<｜Assistant｜>"
+assistant_msg_template: str = "{reasoning}{content}{tool_calls}<｜end▁of▁sentence｜>"
+thinking_template = "{reasoning_content}"
+
+response_format_template: str = (
+    "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}"
+)
+tool_call_template: str = (
+    '<{dsml_token}invoke name="{name}">\n{arguments}\n</{dsml_token}invoke>'
+)
+tool_calls_template = (
+    "<{dsml_token}function_calls>\n{tool_calls}\n</{dsml_token}function_calls>"
+)
+
+tool_output_template: str = "\n<result>{content}</result>"
+
+
+def to_json(value: Any) -> str:
+    try:
+        return json.dumps(value, ensure_ascii=False)
+    except:
+        return json.dumps(value, ensure_ascii=True)
+
+
+def tools_from_openai_format(tools):
+    return [tool["function"] for tool in tools]
+
+
+def tool_calls_from_openai_format(tool_calls):
+    return [
+        {
+            "name": tool_call["function"]["name"],
+            "arguments": tool_call["function"]["arguments"],
+        }
+        for tool_call in tool_calls
+    ]
+
+
+def encode_arguments_to_dsml(tool_call: Dict[str, str]) -> str:
+    p_dsml_template = """<{dsml_token}parameter name="{key}" string="{is_str}">{value}</{dsml_token}parameter>"""
+    P_dsml_strs = []
+
+    arguments = json.loads(tool_call["arguments"])
+
+    for k, v in arguments.items():
+        p_dsml_str = p_dsml_template.format(
+            dsml_token=dsml_token,
+            key=k,
+            is_str="true" if isinstance(v, str) else "false",
+            value=v if isinstance(v, str) else to_json(v),
+        )
+
+        P_dsml_strs.append(p_dsml_str)
+
+    return "\n".join(P_dsml_strs)
+
+
+def decode_dsml_to_arguments(
+    tool_name: str, tool_args: Dict[str, Tuple[str, str]]
+) -> Dict[str, str]:
+    def _decode_value(key: str, value: str, string: str):
+        if string == "true":
+            value = to_json(value)
+        return f"{to_json(key)}: {value}"
+
+    tool_args_json = (
+        "{"
+        + ", ".join(
+            [_decode_value(k, v, string=is_str) for k, (v, is_str) in tool_args.items()]
+        )
+        + "}"
+    )
+    return dict(name=tool_name, arguments=tool_args_json)
+
+
+def render_tools(tools: List[Dict[str, Union[str, Dict[str, Any]]]]) -> str:
+    tools_json = [to_json(t) for t in tools]
+
+    return TOOLS_SYSTEM_TEMPLATE.format(
+        tool_schemas="\n".join(tools_json),
+        dsml_token=dsml_token,
+        thinking_start_token=thinking_start_token,
+        thinking_end_token=thinking_end_token,
+    )
+
+
+def find_last_user_index(messages: List[Dict[str, Any]]) -> int:
+    last_user_index = -1
+    for idx in range(len(messages) - 1, -1, -1):
+        if messages[idx].get("role") in ["user", "developer"]:
+            last_user_index = idx
+            break
+    return last_user_index
+
+
+def render_message(
+    index: int, messages: List[Dict[str, Any]], thinking_mode: str
+) -> str:
+    assert 0 <= index < len(messages)
+    assert thinking_mode in [
+        "chat",
+        "thinking",
+    ], f"Invalid thinking_mode `{thinking_mode}`"
+
+    prompt = ""
+    msg = messages[index]
+    last_user_idx = find_last_user_index(messages)
+
+    role = msg.get("role")
+    content = msg.get("content")
+    tools = msg.get("tools")
+    response_format = msg.get("response_format")
+    tool_calls = msg.get("tool_calls")
+    reasoning_content = msg.get("reasoning_content")
+
+    if tools:
+        tools = tools_from_openai_format(tools)
+    if tool_calls:
+        tool_calls = tool_calls_from_openai_format(tool_calls)
+
+    if role == "system":
+        prompt += system_msg_template.format(content=content or "")
+        if tools:
+            prompt += "\n\n" + render_tools(tools)
+
+        if response_format:
+            prompt += "\n\n" + response_format_template.format(
+                schema=to_json(response_format)
+            )
+
+    elif role == "developer":
+        assert content, f"Invalid message for role `{role}`: {msg}"
+        content_developer = ""
+        if tools:
+            content_developer += "\n\n" + render_tools(tools)
+
+        if response_format:
+            content_developer += "\n\n" + response_format_template.format(
+                schema=to_json(response_format)
+            )
+
+        content_developer += "\n\n# The user's message is: {}".format(content)
+
+        prompt += user_msg_template.format(content=content_developer)
+        if index == last_user_idx and thinking_mode == "thinking":
+            prompt += thinking_start_token
+        else:
+            prompt += thinking_end_token
+
+    elif role == "user":
+        prompt += user_msg_template.format(content=content)
+
+        if index == last_user_idx and thinking_mode == "thinking":
+            prompt += thinking_start_token
+        else:
+            prompt += thinking_end_token
+
+    elif role == "tool":
+        prev_assistant_idx = index - 1
+        assistant_msg = messages[prev_assistant_idx]
+        while prev_assistant_idx >= 0 and assistant_msg.get("role") == "tool":
+            prev_assistant_idx -= 1
+            assistant_msg = messages[prev_assistant_idx]
+
+        assert (
+            index == 0
+            or prev_assistant_idx >= 0
+            and assistant_msg.get("role") == "assistant"
+        ), f"Invalid messages at {index}:\n{assistant_msg}"
+
+        tool_call_order = index - prev_assistant_idx
+        assistant_tool_calls = assistant_msg.get("tool_calls")
+        assert (
+            assistant_tool_calls and len(assistant_tool_calls) >= tool_call_order
+        ), "No tool calls but found tool output"
+
+        if tool_call_order == 1:
+            prompt += "\n\n<function_results>"
+
+        prompt += tool_output_template.format(content=content)
+
+        if tool_call_order == len(assistant_tool_calls):
+            prompt += "\n</function_results>"
+
+            if index >= last_user_idx and thinking_mode == "thinking":
+                prompt += "\n\n" + thinking_start_token
+            else:
+                prompt += "\n\n" + thinking_end_token
+
+    elif role == "assistant":
+        prev_assistant_idx = index
+        thinking_part = ""
+
+        tool_calls_content = ""
+        if tool_calls:
+            tool_calls = [
+                tool_call_template.format(
+                    dsml_token=dsml_token,
+                    name=tool_call.get("name"),
+                    arguments=encode_arguments_to_dsml(tool_call),
+                )
+                for tool_call in tool_calls
+            ]
+            tool_calls_content += "\n\n" + tool_calls_template.format(
+                dsml_token=dsml_token, tool_calls="\n".join(tool_calls)
+            )
+
+        summary_content = content or ""
+
+        if thinking_mode == "thinking" and index > last_user_idx:
+            assert (
+                reasoning_content or tool_calls
+            ), f"ThinkingMode: {thinking_mode}, invalid message without reasoning_content/tool_calls `{msg}` after last user message"
+            thinking_part = (
+                thinking_template.format(reasoning_content=reasoning_content or "")
+                + thinking_end_token
+            )
+
+        prompt += assistant_msg_template.format(
+            reasoning=thinking_part,
+            content=summary_content,
+            tool_calls=tool_calls_content,
+        )
+    else:
+        raise NotImplementedError(f"Unknown role: {role}")
+
+    return prompt
+
+
+def drop_thinking_messages(
+    messages: List[Dict[str, Any]], last_user_idx: Optional[int] = None
+) -> List[Dict[str, Any]]:
+    messages_wo_thinking: List[Dict[str, Any]] = []
+    last_user_idx = (
+        find_last_user_index(messages) if last_user_idx is None else last_user_idx
+    )
+    for idx, msg in enumerate(messages):
+        role = msg.get("role")
+        if role in ["user", "system", "tool"] or idx >= last_user_idx:
+            messages_wo_thinking.append(msg)
+            continue
+
+        elif role == "assistant":
+            msg_wo_thinking = copy.copy(msg)
+            msg_wo_thinking.pop("reasoning_content", None)
+            messages_wo_thinking.append(msg_wo_thinking)
+
+    return messages_wo_thinking
+
+
+def encode_messages(
+    messages: List[Dict[str, Any]],
+    thinking_mode: str = "thinking",
+    context: Optional[List[Dict[str, Any]]] = None,
+    drop_thinking: bool = True,
+    add_default_bos_token: bool = True,
+) -> str:
+    context = context if context else []
+    full_messages = context + messages
+    prompt = bos_token if add_default_bos_token and len(context) == 0 else ""
+
+    if thinking_mode == "thinking" and drop_thinking:
+        full_messages = drop_thinking_messages(full_messages)
+
+    for idx in range(len(messages)):
+        prompt += render_message(
+            idx + len(context), full_messages, thinking_mode=thinking_mode
+        )
+
+    return prompt
+
+
+def apply_chat_template(
+    messages, continue_final_message=False, add_generation_prompt=False, **kwargs
+):
+    out = encode_messages(messages, **kwargs)
+    if continue_final_message and add_generation_prompt:
+        raise ValueError(
+            "Only one of continue_final_message or add_generation_prompt can be True"
+        )
+    if not add_generation_prompt and messages[-1]["role"] == "user":
+        out = out.removesuffix("<｜Assistant｜><think>")
+    if continue_final_message and messages[-1]["role"] == "assistant":
+        out = out.removesuffix(eos_token)
+    return out

mlx_lm/convert.py

@@ -10,8 +10,7 @@ from mlx.utils import tree_map_with_path
 
 from .utils import (
     dequantize_model,
-    fetch_from_hub,
-    get_model_path,
+    load,
     quantize_model,
     save,
     upload_to_hub,
@@ -63,7 +62,9 @@ def mixed_quant_predicate_builder(
             or index >= 7 * num_layers // 8
             or (index - num_layers // 8) % 3 == 2
         )
-        if "v_proj" in path and use_more_bits:
+        if (
+            "v_proj" in path or "v_a_proj" in path or "v_b_proj" in path
+        ) and use_more_bits:
             return {"group_size": group_size, "bits": high_bits}
         if "down_proj" in path and use_more_bits:
             return {"group_size": group_size, "bits": high_bits}
@@ -107,9 +108,12 @@ def convert(
         )
 
     print("[INFO] Loading")
-    model_path, hf_path = get_model_path(hf_path, revision=revision)
-    model, config, tokenizer = fetch_from_hub(
-        model_path, lazy=True, trust_remote_code=trust_remote_code
+    model, tokenizer, config = load(
+        hf_path,
+        revision=revision,
+        return_config=True,
+        tokenizer_config={"trust_remote_code": trust_remote_code},
+        lazy=True,
     )
 
     if isinstance(quant_predicate, str):
@@ -154,11 +158,10 @@ def convert(
 
     save(
         mlx_path,
-        model_path,
+        hf_path,
         model,
         tokenizer,
         config,
-        hf_repo=hf_path,
     )
 
     if upload_repo is not None:
@@ -176,7 +179,12 @@ def configure_parser() -> argparse.ArgumentParser:
         description="Convert Hugging Face model to MLX format"
     )
 
-    parser.add_argument("--hf-path", type=str, help="Path to the Hugging Face model.")
+    parser.add_argument(
+        "--hf-path",
+        "--model",
+        type=str,
+        help="Path to the model. This can be a local path or a Hugging Face Hub model identifier.",
+    )
     parser.add_argument(
         "--mlx-path", type=str, default="mlx_model", help="Path to save the MLX model."
     )
@@ -184,17 +192,23 @@ def configure_parser() -> argparse.ArgumentParser:
         "-q", "--quantize", help="Generate a quantized model.", action="store_true"
     )
     parser.add_argument(
-        "--q-group-size", help="Group size for quantization.", type=int, default=64
+        "--q-group-size",
+        help="Group size for quantization.",
+        type=int,
+        default=None,
     )
     parser.add_argument(
-        "--q-bits", help="Bits per weight for quantization.", type=int, default=4
+        "--q-bits",
+        help="Bits per weight for quantization.",
+        type=int,
+        default=None,
     )
     parser.add_argument(
         "--q-mode",
         help="The quantization mode.",
         type=str,
         default="affine",
-        choices=["affine", "mxfp4"],
+        choices=["affine", "mxfp4", "nvfp4", "mxfp8"],
     )
     parser.add_argument(
         "--quant-predicate",

mlx_lm/evaluate.py

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

mlx_lm/examples/batch_generate_response.py

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

mlx_lm/examples/chat.py

@@ -15,7 +15,10 @@ 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)
+prompt = tokenizer.apply_chat_template(
+    messages,
+    add_generation_prompt=True,
+)
 
 # Assistant response
 response = generate(
@@ -29,7 +32,10 @@ response = generate(
 # User turn
 prompt = "What's my name?"
 messages = [{"role": "user", "content": prompt}]
-prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
+prompt = tokenizer.apply_chat_template(
+    messages,
+    add_generation_prompt=True,
+)
 
 # Assistant response
 response = generate(

mlx_lm/examples/generate_response.py

@@ -14,7 +14,8 @@ conversation = [{"role": "user", "content": prompt}]
 
 # Transform the prompt into the chat template
 prompt = tokenizer.apply_chat_template(
-    conversation=conversation, add_generation_prompt=True
+    conversation=conversation,
+    add_generation_prompt=True,
 )
 
 # Specify the maximum number of tokens

mlx_lm/examples/lora_config.yaml

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

mlx_lm/examples/openai_reasoning_content.py

@@ -0,0 +1,40 @@
+from openai import OpenAI
+
+client = OpenAI(
+    api_key="not-needed",
+    base_url="http://localhost:8080/v1",
+)
+
+model = "mlx-community/Qwen3-4B-Thinking-2507-4bit"
+
+messages = [{"role": "user", "content": "9.11 and 9.8, which is greater?"}]
+
+# Non-streaming example
+
+response = client.chat.completions.create(
+    model=model, messages=messages, max_tokens=2048
+)
+
+reasoning = response.choices[0].message.reasoning
+content = response.choices[0].message.content
+
+print("=== reasoning ===\n")
+print(f"\033[37m{reasoning}\033[0m")
+print("=== content ===\n")
+print(content)
+
+# Streaming example
+
+stream = client.chat.completions.create(
+    model=model,
+    messages=messages,
+    stream=True,
+    max_tokens=2048,
+)
+
+for chunk in stream:
+    if (reasoning := chunk.choices[0].delta.reasoning) is not None:
+        print(f"\033[37m{reasoning}\033[0m", end="")
+    if (content := chunk.choices[0].delta.content) is not None:
+        print(f"{content}", end="")
+print()

mlx_lm/examples/openai_tool_use.py

@@ -8,11 +8,13 @@ To run, first start the server:
 
 Then run this script.
 """
+import json
+
 from openai import OpenAI
 
 client = OpenAI(base_url="http://localhost:8080/v1", api_key="not-needed")
 
-model = "mlx-community/qwen3-4b-4bit-DWQ"
+model = "mlx-community/Qwen3-4B-Instruct-2507-4bit"
 messages = [{"role": "user", "content": "What's the weather in Boston?"}]
 
 tools = [

mlx_lm/examples/pipeline_generate.py

@@ -1,135 +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 json
-import resource
-from pathlib import Path
-
-import mlx.core as mx
-from huggingface_hub import snapshot_download
-from mlx.utils import tree_flatten
-
-from mlx_lm import load, stream_generate
-from mlx_lm.utils import load_model, load_tokenizer
-
-# Needed for 8 bit model
-resource.setrlimit(resource.RLIMIT_NOFILE, (2048, 4096))
-
-
-def download(repo: str, allow_patterns: list[str]) -> Path:
-    return Path(
-        snapshot_download(
-            repo,
-            allow_patterns=allow_patterns,
-        )
-    )
-
-
-def shard_and_load(repo):
-    # Get model path with everything but weight safetensors
-    model_path = download(
-        args.model,
-        allow_patterns=["*.json", "*.py", "tokenizer.model", "*.tiktoken", "*.txt"],
-    )
-
-    # Lazy load and shard model to figure out
-    # which weights we need
-    model, config = load_model(model_path, lazy=True, strict=False)
-
-    group = mx.distributed.init()
-    rank = group.rank()
-    model.model.pipeline(group)
-
-    # Figure out which files we need for the local shard
-    with open(model_path / "model.safetensors.index.json", "r") as fid:
-        weight_index = json.load(fid)["weight_map"]
-
-    local_files = set()
-    for k, _ in tree_flatten(model.parameters()):
-        local_files.add(weight_index[k])
-
-    # Download weights for local shard
-    download(args.model, allow_patterns=local_files)
-
-    # Load and shard the model, and load the weights
-    tokenizer = load_tokenizer(
-        model_path,
-        {"trust_remote_code": True},
-        eos_token_ids=config.get("eos_token_id", None),
-    )
-    model, _ = load_model(model_path, lazy=True, strict=False)
-    model.model.pipeline(group)
-    mx.eval(model.parameters())
-
-    # Synchronize processes before generation to avoid timeout if downloading
-    # model for the first time.
-    mx.eval(mx.distributed.all_sum(mx.array(1.0), stream=mx.cpu))
-    return model, tokenizer
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="LLM pipelined inference example")
-    parser.add_argument(
-        "--model",
-        default="mlx-community/DeepSeek-R1-3bit",
-        help="HF repo or path to local model.",
-    )
-    parser.add_argument(
-        "--prompt",
-        "-p",
-        default="Write a quicksort in C++.",
-        help="Message to be processed by the model ('-' reads from stdin)",
-    )
-    parser.add_argument(
-        "--max-tokens",
-        "-m",
-        type=int,
-        default=256,
-        help="Maximum number of tokens to generate",
-    )
-    args = parser.parse_args()
-
-    group = mx.distributed.init()
-    rank = group.rank()
-
-    def rprint(*args, **kwargs):
-        if rank == 0:
-            print(*args, **kwargs)
-
-    model, tokenizer = shard_and_load(args.model)
-
-    messages = [{"role": "user", "content": args.prompt}]
-    prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True)
-
-    for response in stream_generate(
-        model, tokenizer, prompt, max_tokens=args.max_tokens
-    ):
-        rprint(response.text, end="", flush=True)
-
-    rprint()
-    rprint("=" * 10)
-    rprint(
-        f"Prompt: {response.prompt_tokens} tokens, "
-        f"{response.prompt_tps:.3f} tokens-per-sec"
-    )
-    rprint(
-        f"Generation: {response.generation_tokens} tokens, "
-        f"{response.generation_tps:.3f} tokens-per-sec"
-    )
-    rprint(f"Peak memory: {response.peak_memory:.3f} GB")

mlx_lm/examples/sharded_generate.py

@@ -0,0 +1,86 @@
+# Copyright © 2025 Apple Inc.
+
+"""
+Run with:
+
+```
+mlx.launch \
+    --backend jaccl \
+    --env MLX_METAL_FAST_SYNCH=1 \
+    --hostfile /path/to/hosts.json \
+    /path/to/sharded_generate.py \
+    --prompt 'Hello world'
+```
+
+For more information on running distributed programs with MLX 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 sharded_load
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="LLM distributed inference example")
+    parser.add_argument(
+        "--model",
+        default="mlx-community/Llama-3.3-70B-Instruct-4bit",
+        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",
+    )
+    parser.add_argument(
+        "--pipeline",
+        action="store_true",
+        help="Use pipelining instead of tensor parallelism",
+    )
+    args = parser.parse_args()
+
+    group = mx.distributed.init()
+    rank = group.rank()
+    pipeline_group = group if args.pipeline else None
+    tensor_group = group if not args.pipeline else None
+
+    def rprint(*args, **kwargs):
+        if rank == 0:
+            print(*args, **kwargs)
+
+    model, tokenizer = sharded_load(args.model, pipeline_group, tensor_group)
+
+    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

@@ -6,7 +6,7 @@ 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"
+checkpoint = "mlx-community/Qwen3-4B-Instruct-2507-4bit"
 
 # Load the corresponding model and tokenizer
 model, tokenizer = load(path_or_hf_repo=checkpoint)
@@ -31,7 +31,9 @@ prompt = "Multiply 12234585 and 48838483920."
 messages = [{"role": "user", "content": prompt}]
 
 prompt = tokenizer.apply_chat_template(
-    messages, add_generation_prompt=True, tools=list(tools.values())
+    messages,
+    add_generation_prompt=True,
+    tools=list(tools.values()),
 )
 
 prompt_cache = make_prompt_cache(model)
@@ -47,12 +49,11 @@ response = generate(
 )
 
 # 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())
+# - The tool call format is model specific.
+# - The tokenizer's tool parser expects tool call text to be already extracted.
+start_tool = response.find(tokenizer.tool_call_start) + len(tokenizer.tool_call_start)
+end_tool = response.find(tokenizer.tool_call_end)
+tool_call = tokenizer.tool_parser(response[start_tool:end_tool].strip())
 tool_result = tools[tool_call["name"]](**tool_call["arguments"])
 
 # Put the tool result in the prompt

mlx_lm/fuse.py

@@ -4,10 +4,9 @@ from pathlib import Path
 from mlx.utils import tree_flatten, tree_unflatten
 
 from .gguf import convert_to_gguf
-from .tuner.utils import dequantize, load_adapters
 from .utils import (
-    fetch_from_hub,
-    get_model_path,
+    dequantize_model,
+    load,
     save,
     upload_to_hub,
 )
@@ -40,8 +39,8 @@ def parse_arguments() -> argparse.Namespace:
         default=None,
     )
     parser.add_argument(
-        "--de-quantize",
-        help="Generate a de-quantized model.",
+        "--dequantize",
+        help="Generate a dequantized model.",
         action="store_true",
     )
     parser.add_argument(
@@ -62,14 +61,12 @@ def main() -> None:
     print("Loading pretrained model")
     args = parse_arguments()
 
-    model_path, hf_path = get_model_path(args.model)
-    model, config, tokenizer = fetch_from_hub(model_path)
-
-    model.freeze()
-    model = load_adapters(model, args.adapter_path)
+    model, tokenizer, config = load(
+        args.model, adapter_path=args.adapter_path, return_config=True
+    )
 
     fused_linears = [
-        (n, m.fuse(de_quantize=args.de_quantize))
+        (n, m.fuse(dequantize=args.dequantize))
         for n, m in model.named_modules()
         if hasattr(m, "fuse")
     ]
@@ -77,19 +74,19 @@ def main() -> None:
     if fused_linears:
         model.update_modules(tree_unflatten(fused_linears))
 
-    if args.de_quantize:
-        print("De-quantizing model")
-        model = dequantize(model)
+    if args.dequantize:
+        print("Dequantizing model")
+        model = dequantize_model(model)
         config.pop("quantization", None)
+        config.pop("quantization_config", None)
 
     save_path = Path(args.save_path)
     save(
         save_path,
-        model_path,
+        args.model,
         model,
         tokenizer,
         config,
-        hf_repo=hf_path,
         donate_model=False,
     )
 
@@ -100,13 +97,9 @@ def main() -> None:
                 f"Model type {model_type} not supported for GGUF conversion."
             )
         weights = dict(tree_flatten(model.parameters()))
-        convert_to_gguf(model_path, weights, config, str(save_path / args.gguf_path))
+        convert_to_gguf(save_path, weights, config, str(save_path / args.gguf_path))
 
     if args.upload_repo is not None:
-        if hf_path is None:
-            raise ValueError(
-                "Must provide original Hugging Face repo to upload local model."
-            )
         upload_to_hub(args.save_path, args.upload_repo)
 
 

mlx_lm/generate.py

@@ -7,6 +7,7 @@ import json
 import sys
 import time
 from dataclasses import dataclass
+from functools import partial
 from typing import (
     Any,
     Callable,
@@ -26,8 +27,11 @@ from .models import cache
 from .models.cache import (
     ArraysCache,
     BatchKVCache,
+    BatchRotatingKVCache,
+    CacheList,
     KVCache,
     QuantizedKVCache,
+    RotatingKVCache,
     load_prompt_cache,
 )
 from .sample_utils import make_sampler
@@ -177,8 +181,7 @@ def setup_arg_parser():
     parser.add_argument(
         "--kv-bits",
         type=int,
-        help="Number of bits for KV cache quantization. "
-        "Defaults to no quantization.",
+        help="Number of bits for KV cache quantization. Defaults to no quantization.",
         default=None,
     )
     parser.add_argument(
@@ -284,16 +287,11 @@ class GenerationResponse:
 
 
 def maybe_quantize_kv_cache(prompt_cache, quantized_kv_start, kv_group_size, kv_bits):
-    if (
-        kv_bits is not None
-        and not isinstance(prompt_cache[0], cache.QuantizedKVCache)
-        and prompt_cache[0].offset > quantized_kv_start
-    ):
-        for i in range(len(prompt_cache)):
-            if isinstance(prompt_cache[i], cache.KVCache):
-                prompt_cache[i] = prompt_cache[i].to_quantized(
-                    group_size=kv_group_size, bits=kv_bits
-                )
+    if kv_bits is None:
+        return
+    for e, c in enumerate(prompt_cache):
+        if hasattr(c, "to_quantized") and c.offset >= quantized_kv_start:
+            prompt_cache[e] = c.to_quantized(group_size=kv_group_size, bits=kv_bits)
 
 
 def generate_step(
@@ -301,7 +299,7 @@ def generate_step(
     model: nn.Module,
     *,
     max_tokens: int = 256,
-    sampler: Optional[Callable[mx.array, mx.array]] = None,
+    sampler: Optional[Callable[[mx.array], mx.array]] = None,
     logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None,
     max_kv_size: Optional[int] = None,
     prompt_cache: Optional[Any] = None,
@@ -309,7 +307,7 @@ def generate_step(
     kv_bits: Optional[int] = None,
     kv_group_size: int = 64,
     quantized_kv_start: int = 0,
-    prompt_progress_callback: Optional[Callable[int, int]] = None,
+    prompt_progress_callback: Optional[Callable[[int, int], None]] = None,
     input_embeddings: Optional[mx.array] = None,
 ) -> Generator[Tuple[mx.array, mx.array], None, None]:
     """
@@ -335,7 +333,7 @@ def generate_step(
         kv_group_size (int): Group size for KV cache quantization. Default: ``64``.
         quantized_kv_start (int): Step to begin using a quantized KV cache.
            when ``kv_bits`` is non-None. Default: ``0``.
-        prompt_progress_callback (Callable[int, int]): A call-back which takes the
+        prompt_progress_callback (Callable[[int, int], None]): A call-back which takes the
            prompt tokens processed so far and the total number of prompt tokens.
         input_embeddings (mx.array, optional): Input embeddings to use instead of or in
           conjunction with prompt tokens. Default: ``None``.
@@ -420,7 +418,8 @@ def generate_step(
         prompt_processed_tokens = 0
         prompt_progress_callback(prompt_processed_tokens, total_prompt_tokens)
         while total_prompt_tokens - prompt_processed_tokens > 1:
-            n_to_process = min(prefill_step_size, prompt.size - 1)
+            remaining = (total_prompt_tokens - prompt_processed_tokens) - 1
+            n_to_process = min(prefill_step_size, remaining)
             _model_call(
                 input_tokens=prompt[:n_to_process][None],
                 input_embeddings=(
@@ -468,7 +467,7 @@ def speculative_generate_step(
     *,
     num_draft_tokens: int = 2,
     max_tokens: int = 256,
-    sampler: Optional[Callable[mx.array, mx.array]] = None,
+    sampler: Optional[Callable[[mx.array], mx.array]] = None,
     logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None,
     prompt_cache: Optional[Any] = None,
     prefill_step_size: int = 512,
@@ -487,7 +486,7 @@ def speculative_generate_step(
           speculative decoding. Default: ``2``.
         max_tokens (int): The maximum number of tokens. Use``-1`` for an infinite
           generator. Default: ``256``.
-        sampler (Callable[mx.array, mx.array], optional): A sampler for sampling a
+        sampler (Callable[[mx.array], mx.array], optional): A sampler for sampling a
           token from a vector of log probabilities. Default: ``None``.
         logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
           A list of functions that take tokens and logits and return the processed
@@ -548,7 +547,9 @@ def speculative_generate_step(
                     y = y[: -(n_predict - 1)]
                 for i in range(n_predict):
                     prev_tokens = (
-                        mx.concat([prev_tokens, y]) if prev_tokens is not None else y
+                        mx.concatenate([prev_tokens, y])
+                        if prev_tokens is not None
+                        else y
                     )
                     y, logprobs = _process_and_sample(prev_tokens, logits[:, i, :])
                     out_y.append(y)
@@ -787,6 +788,12 @@ def _left_pad_prompts(prompts, max_length=None):
     return mx.array([[0] * (max_length - len(p)) + p for p in prompts])
 
 
+def _right_pad_prompts(prompts, max_length=None):
+    if max_length is None:
+        max_length = max(len(p) for p in prompts)
+    return mx.array([p + [0] * (max_length - len(p)) for p in prompts])
+
+
 @dataclass
 class BatchStats:
     """
@@ -823,6 +830,7 @@ class BatchResponse:
 
     texts: List[str]
     stats: BatchStats
+    caches: Optional[List[List[Any]]]
 
 
 @dataclass
@@ -833,129 +841,306 @@ class Batch:
     max_tokens: List[int]
     num_tokens: List[int]
     cache: List[Any]
+    samplers: List[Any]
+    logits_processors: List[Any]
+    tokens: List[mx.array]
 
     def __len__(self):
         return len(self.uids)
 
     def filter(self, keep_idx: List[int]):
         self.uids = [self.uids[k] for k in keep_idx]
+        self.logprobs = [self.logprobs[k] for k in keep_idx]
         self.max_tokens = [self.max_tokens[k] for k in keep_idx]
         self.num_tokens = [self.num_tokens[k] for k in keep_idx]
+        self.samplers = [self.samplers[k] for k in keep_idx]
+        self.logits_processors = [self.logits_processors[k] for k in keep_idx]
+        self.tokens = [self.tokens[k] for k in keep_idx]
         keep_idx = mx.array(keep_idx, mx.int32)
         self.y = self.y[keep_idx]
-        self.logprobs = self.logprobs[keep_idx]
         for c in self.cache:
             c.filter(keep_idx)
 
     def extend(self, other):
         self.uids.extend(other.uids)
         self.y = mx.concatenate([self.y, other.y])
-        self.logprobs = mx.concatenate([self.logprobs, other.logprobs])
+        self.logprobs.extend(other.logprobs)
         self.num_tokens.extend(other.num_tokens)
         self.max_tokens.extend(other.max_tokens)
+        self.samplers.extend(other.samplers)
+        self.logits_processors.extend(other.logits_processors)
+        self.tokens.extend(other.tokens)
         for c, o in zip(self.cache, other.cache):
             c.extend(o)
 
+    def extract_cache(self, idx):
+        return [c.extract(idx) for c in self.cache]
+
 
 def _make_cache(model, left_padding):
     """
     Convert a list of regular caches into their corresponding
     batch-aware caches.
     """
+
+    def to_batch_cache(c):
+        if type(c) is KVCache:
+            return BatchKVCache(left_padding)
+        elif isinstance(c, ArraysCache):
+            c.left_padding = mx.array(left_padding)
+            return c
+        elif isinstance(c, RotatingKVCache):
+            if c.keep > 0:
+                raise ValueError("RotatingKVCache with keep tokens is not supported.")
+            return BatchRotatingKVCache(c.max_size, left_padding)
+        elif isinstance(c, CacheList):
+            return CacheList(*(to_batch_cache(sub_c) for sub_c in c.caches))
+        else:
+            raise ValueError(f"{type(c)} does not yet support batching")
+
     if hasattr(model, "make_cache"):
         cache = model.make_cache()
-        batch_cache = []
-        for c in cache:
-            if isinstance(c, KVCache):
-                batch_cache.append(BatchKVCache(left_padding))
-            elif isinstance(c, ArraysCache):
-                c.left_padding = mx.array(left_padding)
-                batch_cache.append(c)
-            else:
-                raise ValueError(f"{type(c)} does not yet support batching")
-        return batch_cache
+        return [to_batch_cache(c) for c in cache]
     else:
         return [BatchKVCache(left_padding) for _ in model.layers]
 
 
-class BatchGenerator:
+def _merge_caches(caches):
+    batch_cache = []
+    for i in range(len(caches[0])):
+        if hasattr(caches[0][i], "merge"):
+            batch_cache.append(caches[0][i].merge([c[i] for c in caches]))
+        else:
+            raise ValueError(
+                f"{type(caches[0][i])} does not yet support batching with history"
+            )
+    return batch_cache
 
+
+class BatchGenerator:
     @dataclass
     class Response:
         uid: int
         token: int
         logprobs: mx.array
         finish_reason: Optional[str]
+        prompt_cache: Callable[[], List[Any]]
 
     def __init__(
         self,
         model,
         max_tokens: int = 128,
         stop_tokens: Optional[set] = None,
-        sampler: Optional[Callable[mx.array, mx.array]] = None,
+        sampler: Optional[Callable[[mx.array], mx.array]] = None,
+        logits_processors: Optional[
+            List[Callable[[mx.array, mx.array], mx.array]]
+        ] = None,
         completion_batch_size: int = 32,
         prefill_batch_size: int = 8,
         prefill_step_size: int = 2048,
+        prompt_progress_callback: Optional[
+            Callable[[List[Tuple[int, int, int]]], None]
+        ] = None,
     ):
         self.model = model
         self.unprocessed_prompts = []
         self.max_tokens = max_tokens
         self.stop_tokens = stop_tokens or set()
         self.sampler = sampler or (lambda x: mx.argmax(x, axis=-1))
+        self.logits_processors = logits_processors or []
         self.uid_count = 0
         self.prefill_step_size = prefill_step_size
         self.prefill_batch_size = prefill_batch_size
-        self.completion_batch_size = completion_batch_size
+        self.completion_batch_size = max(completion_batch_size, prefill_batch_size)
+        self.prompt_progress_callback = prompt_progress_callback or (lambda *_: None)
         self._stats = BatchStats()
 
         self.active_batch = None
 
-    def insert(self, prompts, max_tokens: Union[List[int], int, None] = None):
+        if mx.metal.is_available():
+            self._old_wired_limit = mx.set_wired_limit(
+                mx.metal.device_info()["max_recommended_working_set_size"]
+            )
+        else:
+            self._old_wired_limit = None
+
+    def close(self):
+        if self._old_wired_limit is not None:
+            mx.synchronize(generation_stream)
+            mx.set_wired_limit(self._old_wired_limit)
+            self._old_wired_limit = None
+
+    def __del__(self):
+        self.close()
+
+    def insert(
+        self,
+        prompts,
+        max_tokens: Union[List[int], int, None] = None,
+        caches=None,
+        samplers: list | None = None,
+        logits_processors: list | None = None,
+    ):
         uids = []
 
         if max_tokens is None or isinstance(max_tokens, int):
             max_tokens = [max_tokens or self.max_tokens] * len(prompts)
 
-        for p, m in zip(prompts, max_tokens):
-            self.unprocessed_prompts.append((self.uid_count, p, m))
+        if caches is None:
+            caches = [None] * len(prompts)
+        for i in range(len(prompts)):
+            if caches[i] is None:
+                caches[i] = cache.make_prompt_cache(self.model)
+
+        samplers = samplers or [None] * len(prompts)
+        logits_processors = logits_processors or [self.logits_processors] * len(prompts)
+
+        for p, m, c, s, lp in zip(
+            prompts, max_tokens, caches, samplers, logits_processors
+        ):
+            self.unprocessed_prompts.append((self.uid_count, p, m, c, s, lp))
             uids.append(self.uid_count)
             self.uid_count += 1
         # Sort in ascending order of length
         self.unprocessed_prompts = sorted(
-            self.unprocessed_prompts, key=lambda x: len(x[1])
+            self.unprocessed_prompts,
+            key=lambda x: len(x[1]) + max(c.size() for c in x[3]),
         )
         return uids
 
+    def remove(self, uids: List[int]):
+        uids = set(uids)
+        if self.active_batch is not None:
+            batch = self.active_batch
+            keep_idx = [e for e, uid in enumerate(batch.uids) if uid not in uids]
+            if len(keep_idx) > 0:
+                batch.filter(keep_idx)
+            else:
+                self.active_batch = None
+
+        for i in reversed(range(len(self.unprocessed_prompts))):
+            if self.unprocessed_prompts[i][0] in uids:
+                self.unprocessed_prompts.pop(i)
+
     def _process_prompts(self, prompts):
-        uids, inputs, max_tokens = zip(*prompts)
+        uids, inputs, max_tokens, caches, samplers, logits_processors = zip(*prompts)
+
         lengths = [len(p) for p in inputs]
         max_length = max(lengths)
-        batch_size = self.prefill_batch_size
+        padding = [max_length - l for l in lengths]
+
         self._stats.prompt_tokens += sum(lengths)
-        left_padding = [max_length - l for l in lengths]
-        inputs = _left_pad_prompts(inputs, max_length=max_length)
 
-        prompt_cache = _make_cache(self.model, left_padding)
+        tokens = [mx.array(inp) for inp in inputs]
+        processed_tokens = 0
+
+        # New prompts so
+        #   1. Left-pad the inputs
+        #   2. Process
+        if all(c[0].empty() for c in caches):
+            inputs = _left_pad_prompts(inputs, max_length=max_length)
+            prompt_cache = _make_cache(self.model, padding)
+
+            while inputs.shape[1] > 1:
+                n_to_process = min(self.prefill_step_size, inputs.shape[1] - 1)
+                self.model(inputs[:, :n_to_process], cache=prompt_cache)
+                mx.eval([c.state for c in prompt_cache])
+                inputs = inputs[:, n_to_process:]
+                processed_tokens += n_to_process
+                self.prompt_progress_callback(
+                    [
+                        (uid, processed_tokens, length)
+                        for uid, length in zip(uids, lengths)
+                    ]
+                )
+
+        # Further prompt processing so we need to
+        #   1. Merge the KV caches and prepare for right padded prompts
+        #   2. Right pad the inputs
+        #   2. Process
+        #   3. Finalize the KV caches so they are left padded again
+        else:
+            last_inputs = mx.array([p[-1:] for p in inputs])
+            inputs = _right_pad_prompts(inputs, max_length=max_length)
+            prompt_cache = _merge_caches(caches)
+
+            for c in prompt_cache:
+                # subtract one from lengths since we don't process the last token during prefill
+                c.prepare(lengths=[l - 1 for l in lengths], right_padding=padding)
+
+            while inputs.shape[1] > 1:
+                n_to_process = min(self.prefill_step_size, inputs.shape[1] - 1)
+                self.model(inputs[:, :n_to_process], cache=prompt_cache)
+                mx.eval([c.state for c in prompt_cache])
+                inputs = inputs[:, n_to_process:]
+                processed_tokens += n_to_process
+                self.prompt_progress_callback(
+                    [
+                        (uid, processed_tokens, length)
+                        for uid, length in zip(uids, lengths)
+                    ]
+                )
+                mx.clear_cache()
 
-        while inputs.shape[1] > 1:
-            n_to_process = min(self.prefill_step_size, inputs.shape[1] - 1)
-            self.model(inputs[:, :n_to_process], cache=prompt_cache)
             mx.eval([c.state for c in prompt_cache])
-            inputs = inputs[:, n_to_process:]
-            mx.clear_cache()
+            inputs = last_inputs
 
-        y, logprobs = self._step(inputs, prompt_cache)
-        mx.async_eval(y, logprobs)
-        return Batch(
-            list(uids), y, logprobs, list(max_tokens), [0] * len(uids), prompt_cache
+        for c in prompt_cache:
+            c.finalize()
+        mx.clear_cache()
+
+        y, logprobs = self._step(
+            inputs, prompt_cache, samplers, logits_processors, tokens
         )
 
-    def _step(self, input_tokens: mx.array, prompt_cache: List[Any]):
+        mx.async_eval(y, logprobs)
+
+        return Batch(
+            list(uids),
+            y,
+            logprobs,
+            list(max_tokens),
+            [0] * len(uids),
+            prompt_cache,
+            list(samplers),
+            list(logits_processors),
+            tokens,
+        )
+
+    def _step(
+        self,
+        input_tokens: mx.array,
+        prompt_cache: List[Any],
+        samplers: list | None,
+        logits_processors: list | None,
+        tokens: List[mx.array],
+    ):
+        batch_size = input_tokens.shape[0]
+
         logits = self.model(input_tokens, cache=prompt_cache)
         logits = logits[:, -1, :]
+
+        if any(logits_processors):
+            processed_logits = []
+            for e in range(batch_size):
+                sample_logits = logits[e : e + 1]
+                for processor in logits_processors[e]:
+                    sample_logits = processor(tokens[e], sample_logits)
+                processed_logits.append(sample_logits)
+            logits = mx.concatenate(processed_logits, axis=0)
+
         logprobs = logits - mx.logsumexp(logits, axis=-1, keepdims=True)
-        sampled = self.sampler(logprobs)
-        return sampled, logprobs
+        if any(samplers):
+            all_samples = []
+            for e in range(batch_size):
+                sample_sampler = samplers[e] or self.sampler
+                sampled = sample_sampler(logprobs[e : e + 1])
+                all_samples.append(sampled)
+            sampled = mx.concatenate(all_samples, axis=0)
+        else:
+            sampled = self.sampler(logprobs)
+
+        return sampled, list(logprobs)
 
     def stats(self):
         self._stats.prompt_tps = self._stats.prompt_tokens / self._stats.prompt_time
@@ -1004,7 +1189,16 @@ class BatchGenerator:
 
         batch = self.active_batch
         y, logprobs = batch.y, batch.logprobs
-        batch.y, batch.logprobs = self._step(y[:, None], batch.cache)
+        for i, toks in enumerate(batch.tokens):
+            batch.tokens[i] = mx.concatenate((toks, y[i : i + 1]))
+        batch.y, batch.logprobs = self._step(
+            y[:, None],
+            batch.cache,
+            batch.samplers,
+            batch.logits_processors,
+            batch.tokens,
+        )
+
         mx.async_eval(batch.y, batch.logprobs)
 
         y = y.tolist()
@@ -1020,6 +1214,7 @@ class BatchGenerator:
         for e, (t, uid, num_tok, max_tok) in enumerate(
             zip(y, batch.uids, batch.num_tokens, batch.max_tokens)
         ):
+            cache = None
             num_tok += 1
             batch.num_tokens[e] = num_tok
             if t in self.stop_tokens:
@@ -1031,7 +1226,9 @@ class BatchGenerator:
             else:
                 finish_reason = None
                 keep_idx.append(e)
-            responses.append(self.Response(uid, t, logprobs[e], finish_reason))
+            if finish_reason is not None:
+                cache = batch.extract_cache(e)
+            responses.append(self.Response(uid, t, logprobs[e], finish_reason, cache))
 
         # Remove any finished completions
         if len(end_idx):
@@ -1051,9 +1248,12 @@ class BatchGenerator:
 def batch_generate(
     model,
     tokenizer,
-    prompts: List[int],
+    prompts: List[List[int]],
+    prompt_caches: Optional[List[List[Any]]] = None,
     max_tokens: Union[int, List[int]] = 128,
     verbose: bool = False,
+    return_prompt_caches: bool = False,
+    logits_processors: Optional[List[Callable[[mx.array, mx.array], mx.array]]] = None,
     **kwargs,
 ) -> BatchResponse:
     """
@@ -1062,40 +1262,56 @@ def batch_generate(
     Args:
        model (nn.Module): The language model.
        tokenizer (PreTrainedTokenizer): The tokenizer.
-       prompt (List[List[int]]): The input prompts.
+       prompts (List[List[int]]): The input prompts.
+       prompt_caches (List[List[Any]], optional): Pre-computed prompt-caches
+          for each input prompt. Note, unlike ``generate_step``, the caches
+          won't be updated in-place.
        verbose (bool): If ``True``, print tokens and timing information.
           Default: ``False``.
        max_tokens (Union[int, List[int]): Maximum number of output tokens. This
           can be per prompt if a list is provided.
+       return_prompt_caches (bool): Return the prompt caches in the batch
+          responses. Default: ``False``.
+       logits_processors (List[Callable[[mx.array, mx.array], mx.array]], optional):
+          A list of functions that take tokens and logits and return the processed logits. Default: ``None``.
        kwargs: The remaining options get passed to :obj:`BatchGenerator`.
           See :obj:`BatchGenerator` for more details.
     """
 
-    gen = BatchGenerator(model, stop_tokens=tokenizer.eos_token_ids, **kwargs)
+    gen = BatchGenerator(
+        model,
+        stop_tokens=tokenizer.eos_token_ids,
+        **kwargs,
+    )
     num_samples = len(prompts)
     fin = 0
     if verbose:
         print(f"[batch_generate] Finished processing 0/{num_samples} ...", end="\r")
 
-    with wired_limit(model, [generation_stream]):
-        uids = gen.insert(prompts, max_tokens)
-        results = {uid: [] for uid in uids}
-        while responses := gen.next():
-            for r in responses:
-                if verbose and r.finish_reason != None:
+    uids = gen.insert(prompts, max_tokens, caches=prompt_caches)
+    results = {uid: [] for uid in uids}
+    prompt_caches = {}
+    while responses := gen.next():
+        for r in responses:
+            if r.finish_reason is not None:
+                if return_prompt_caches:
+                    prompt_caches[r.uid] = r.prompt_cache
+                if verbose:
                     fin += 1
                     print(
                         f"[batch_generate] Finished processing {fin}/{num_samples} ...",
                         end="\r",
                     )
-                if r.finish_reason != "stop":
-                    results[r.uid].append(r.token)
+            if r.finish_reason != "stop":
+                results[r.uid].append(r.token)
+    gen.close()
     if verbose:
         print(f"[batch_generate] Finished processing {fin}/{num_samples}")
 
     # Return results in correct order
     texts = [tokenizer.decode(results[uid]) for uid in uids]
     stats = gen.stats()
+    caches = [prompt_caches[uid] for uid in uids] if return_prompt_caches else None
     if verbose:
         print(
             f"[batch_generate] Prompt: {stats.prompt_tokens} tokens, {stats.prompt_tps:.3f} tokens-per-sec"
@@ -1105,7 +1321,7 @@ def batch_generate(
             f"{stats.generation_tps:.3f} tokens-per-sec"
         )
         print(f"[batch_generate] Peak memory: {stats.peak_memory:.3f} GB")
-    return BatchResponse(texts, stats)
+    return BatchResponse(texts, stats, caches)
 
 
 def main():
@@ -1162,15 +1378,9 @@ def main():
     if args.chat_template_config is not None:
         template_kwargs = json.loads(args.chat_template_config)
 
-    if args.use_default_chat_template:
-        if tokenizer.chat_template is None:
-            tokenizer.chat_template = tokenizer.default_chat_template
-    elif using_cache:
-        tokenizer.chat_template = json.loads(metadata["chat_template"])
-
     prompt = args.prompt.replace("\\n", "\n").replace("\\t", "\t")
     prompt = sys.stdin.read() if prompt == "-" else prompt
-    if not args.ignore_chat_template and tokenizer.chat_template is not None:
+    if not args.ignore_chat_template and tokenizer.has_chat_template:
         if args.system_prompt is not None:
             messages = [{"role": "system", "content": args.system_prompt}]
         else:

mlx_lm/lora.py

@@ -40,7 +40,7 @@ yaml_loader.add_implicit_resolver(
 )
 
 CONFIG_DEFAULTS = {
-    "model": "mlx_model",
+    "model": "Qwen/Qwen3-0.6b",
     "train": False,
     "fine_tune_type": "lora",
     "optimizer": "adam",
@@ -51,7 +51,7 @@ CONFIG_DEFAULTS = {
         "sgd": {},
         "adafactor": {},
     },
-    "data": "data/",
+    "data": "mlx-community/WikiSQL",
     "seed": 0,
     "num_layers": 16,
     "batch_size": 4,
@@ -68,10 +68,10 @@ CONFIG_DEFAULTS = {
     "max_seq_length": 2048,
     "config": None,
     "grad_checkpoint": False,
+    "grad_accumulation_steps": 1,
     "lr_schedule": None,
     "lora_parameters": {"rank": 8, "dropout": 0.0, "scale": 20.0},
     "mask_prompt": False,
-    "wandb": None,  # will be deprecated in a future release
     "report_to": None,
     "project_name": None,
 }
@@ -142,6 +142,11 @@ def build_parser():
         type=int,
         help="Number of training steps between validations.",
     )
+    parser.add_argument(
+        "--grad-accumulation-steps",
+        type=int,
+        help="Number of steps to accumulate before each optimizer update.",
+    )
     parser.add_argument(
         "--resume-adapter-file",
         type=str,
@@ -185,15 +190,6 @@ def build_parser():
         help="Use gradient checkpointing to reduce memory use.",
         default=None,
     )
-    parser.add_argument(  # will be deprecated in a future release
-        "--wandb",
-        type=str,
-        default=None,
-        help=(
-            "The 'wandb' argument is deprecated and will be removed in a future release. "
-            "Use 'report_to: wandb' and 'project_name' in the configuration instead."
-        ),
-    )
     parser.add_argument(
         "--report-to",
         type=str,
@@ -265,6 +261,7 @@ def train_model(
         adapter_file=adapter_file,
         max_seq_length=args.max_seq_length,
         grad_checkpoint=args.grad_checkpoint,
+        grad_accumulation_steps=args.grad_accumulation_steps,
     )
 
     # Initialize the selected optimizer
@@ -314,14 +311,6 @@ def evaluate_model(args, model: nn.Module, test_set):
 
 def run(args, training_callback: TrainingCallback = None):
     np.random.seed(args.seed)
-    if args.wandb is not None:
-        warnings.warn(
-            "The 'wandb' argument is deprecated and will be removed in a future release. "
-            "Use 'report_to: wandb' and 'project_name' in the configuration instead.",
-            DeprecationWarning,
-        )
-        args.report_to = "wandb"
-        args.project_name = args.wandb
     training_callback = get_reporting_callbacks(
         args.report_to,
         project_name=args.project_name,

mlx_lm/models/Klear.py

@@ -6,6 +6,7 @@ from typing import Any, List, Optional
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .switch_layers import SwitchGLU
 
@@ -114,7 +115,7 @@ class KlearMLP(nn.Module):
         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))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class KlearSparseMoeBlock(nn.Module):

mlx_lm/models/activations.py

@@ -0,0 +1,11 @@
+# Copyright © 2023-2026 Apple Inc.
+
+from functools import partial
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+@partial(mx.compile, shapeless=True)
+def swiglu(gate, x):
+    return nn.silu(gate) * x

mlx_lm/models/afm7.py

@@ -9,6 +9,7 @@ from typing import Any, Dict, Optional, Union
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .cache import ConcatenateKVCache, KVCache
 from .rope_utils import initialize_rope
@@ -50,7 +51,7 @@ class FusedLoRALinear(nn.Module):
         ]
         self.lora_b = [mx.zeros((r, od)) for od in output_dims]
 
-    def fuse(self, de_quantize: bool = False):
+    def fuse(self, dequantize: bool = False):
         linear = self.linear
         weight = linear.weight
         is_quantized = isinstance(linear, FusedQuantizedLinear)
@@ -79,7 +80,7 @@ class FusedLoRALinear(nn.Module):
         delta = mx.concatenate(deltas, axis=0)
         fused_linear.weight = weight + delta
 
-        if is_quantized and not de_quantize:
+        if is_quantized and not dequantize:
             fused_linear = fused_linear.to_quantized(linear.group_size, linear.bits)
 
         return fused_linear
@@ -262,11 +263,6 @@ class KVReuseAttention(nn.Module):
         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__()
@@ -281,7 +277,7 @@ class MLP(nn.Module):
     def __call__(self, x) -> mx.array:
         g = self.gate_proj(x)
         x = self.up_proj(x)
-        return self.down_proj(_swiglu(g, x))
+        return self.down_proj(swiglu(g, x))
 
 
 class TransformerBlock(nn.Module):

mlx_lm/models/afmoe.py

@@ -0,0 +1,405 @@
+# 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 .activations import swiglu
+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(swiglu(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/baichuan_m1.py

@@ -6,6 +6,7 @@ from typing import Any, List, Optional
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .cache import CacheList, KVCache, MambaCache, RotatingKVCache
 
@@ -140,7 +141,7 @@ class MLP(nn.Module):
         )
 
     def __call__(self, x: mx.array) -> mx.array:
-        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class DecoderLayer(nn.Module):

mlx_lm/models/bailing_moe.py

@@ -1,11 +1,13 @@
 # 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 .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .rope_utils import initialize_rope
 from .switch_layers import SwitchGLU
@@ -37,6 +39,7 @@ class ModelArgs(BaseModelArgs):
     use_qk_norm: bool = False
     tie_word_embeddings: bool = False
     partial_rotary_factor: float = 1.0
+    rotary_dim: Optional[int] = None
     moe_router_enable_expert_bias: bool = False
     moe_router_enable_routed_scaling: bool = True
     routed_scaling_factor: float = 1.0
@@ -47,6 +50,13 @@ class ModelArgs(BaseModelArgs):
     moe_router_enable_shared_expert: bool = True
 
 
+@partial(mx.compile, shapeless=True)
+def aggregate_expert_outputs(expert_outputs, scores):
+    return (
+        (expert_outputs * scores[..., None]).sum(axis=-2).astype(expert_outputs.dtype)
+    )
+
+
 class BailingMoeMLP(nn.Module):
     def __init__(self, args: ModelArgs, intermediate_size: Optional[int] = None):
         super().__init__()
@@ -67,7 +77,7 @@ class BailingMoeMLP(nn.Module):
         )
 
     def __call__(self, x) -> mx.array:
-        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class BailingMoeAttention(nn.Module):
@@ -94,8 +104,10 @@ class BailingMoeAttention(nn.Module):
             self.key_layernorm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
             self.query_layernorm = nn.RMSNorm(self.head_dim, eps=args.rms_norm_eps)
 
+        if (rope_dim := args.rotary_dim) is None:
+            rope_dim = int(self.head_dim * args.partial_rotary_factor)
         self.rope = initialize_rope(
-            int(self.head_dim * args.partial_rotary_factor),
+            rope_dim,
             args.rope_theta,
             traditional=False,
             scaling_config=args.rope_scaling,
@@ -146,6 +158,7 @@ class BailingMoeAttention(nn.Module):
         return self.dense(output)
 
 
+@mx.compile
 def group_expert_select(
     gates,
     e_score_correction_bias,
@@ -171,15 +184,15 @@ def group_expert_select(
         k = n_group - topk_group
         group_idx = mx.argpartition(group_scores, kth=k - 1, axis=-2)[..., :k, :]
         scores = mx.put_along_axis(
-            scores, mx.stop_gradient(group_idx), mx.array(0.0), axis=-2
+            scores, mx.stop_gradient(group_idx), mx.array(0.0, scores.dtype), axis=-2
         )
         scores = mx.flatten(scores, -2, -1)
 
     k = top_k
-    inds = mx.argpartition(-scores, kth=k - 1, axis=-1)[..., :k]
+    inds = mx.argpartition(scores, kth=-k, axis=-1)[..., -k:]
     scores = mx.take_along_axis(orig_scores, inds, axis=-1)
     if top_k > 1 and norm_topk_prob:
-        denominator = scores.sum(axis=-1, keepdims=True)
+        denominator = scores.sum(axis=-1, keepdims=True) + 1e-20
         scores = scores / denominator
     scores = scores * routed_scaling_factor
 
@@ -245,7 +258,7 @@ class BailingMoeSparseMoeBlock(nn.Module):
     def __call__(self, x):
         topk_idx, topk_weight = self.gate(x)
         out = self.switch_mlp(x, topk_idx)
-        out = (out * topk_weight[..., None]).sum(axis=-2)
+        out = aggregate_expert_outputs(out, topk_weight)
         if self.shared_experts is not None:
             out = out + self.shared_experts(x)
         return out

mlx_lm/models/bailing_moe_linear.py

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

mlx_lm/models/cache.py

@@ -1,5 +1,6 @@
 # Copyright © 2023-2024 Apple Inc.
 
+import copy
 from typing import Any, Dict, List, Optional
 
 import mlx.core as mx
@@ -73,10 +74,10 @@ def load_prompt_cache(file_name, return_metadata=False):
     arrays = tree_unflatten(list(arrays.items()))
     cache_metadata = tree_unflatten(list(cache_metadata.items()))
     info, metadata, classes = cache_metadata
-    cache = [globals()[c]() for c in classes]
-    for c, state, meta_state in zip(cache, arrays, info):
-        c.state = state
-        c.meta_state = meta_state
+    cache = [
+        globals()[c].from_state(state, meta_state)
+        for c, state, meta_state in zip(classes, arrays, info)
+    ]
     if return_metadata:
         return cache, metadata
     return cache
@@ -141,6 +142,29 @@ class _BaseCache:
     def is_trimmable(self):
         return False
 
+    def size(self):
+        """
+        Return the size (i.e. sequence length) of the cache.
+
+        Not every cache is required to implement this, in which case the size
+        will always be 0 (though the cache may not be empty).
+        """
+        return 0
+
+    def empty(self):
+        """
+        Return if the cache is empty or not.
+        """
+        raise NotImplementedError("Cache sub-class must implement this.")
+
+    @classmethod
+    def from_state(cls, state, meta_state):
+        # Create an instance of cls without calling __init__
+        obj = cls.__new__(cls)
+        obj.state = state
+        obj.meta_state = meta_state
+        return obj
+
 
 class ConcatenateKVCache(_BaseCache):
     """ConcatenateKVCache the simplest KV cache implementation.
@@ -186,13 +210,17 @@ class ConcatenateKVCache(_BaseCache):
     def make_mask(self, *args, **kwargs):
         return create_attention_mask(*args, offset=self.offset, **kwargs)
 
+    def empty(self):
+        return self.keys is None
+
 
 class QuantizedKVCache(_BaseCache):
+    step = 256
+
     def __init__(self, group_size: int = 64, bits: int = 8):
         self.keys = None
         self.values = None
         self.offset = 0
-        self.step = 256
         self.group_size = group_size
         self.bits = bits
 
@@ -254,11 +282,11 @@ class QuantizedKVCache(_BaseCache):
 
     @property
     def meta_state(self):
-        return tuple(map(str, (self.step, self.offset, self.group_size, self.bits)))
+        return tuple(map(str, (self.offset, self.group_size, self.bits)))
 
     @meta_state.setter
     def meta_state(self, v):
-        self.step, self.offset, self.group_size, self.bits = map(int, v)
+        self.offset, self.group_size, self.bits = map(int, v)
 
     def is_trimmable(self):
         return True
@@ -271,13 +299,17 @@ class QuantizedKVCache(_BaseCache):
     def make_mask(self, *args, **kwargs):
         return create_attention_mask(*args, offset=self.offset, **kwargs)
 
+    def empty(self):
+        return self.keys is None
+
 
 class KVCache(_BaseCache):
+    step = 256
+
     def __init__(self):
         self.keys = None
         self.values = None
         self.offset = 0
-        self.step = 256
 
     def update_and_fetch(self, keys, values):
         prev = self.offset
@@ -303,6 +335,9 @@ class KVCache(_BaseCache):
         self.values[..., prev : self.offset, :] = values
         return self.keys[..., : self.offset, :], self.values[..., : self.offset, :]
 
+    def size(self):
+        return self.offset
+
     @property
     def state(self):
         if self.offset == self.keys.shape[2]:
@@ -339,16 +374,23 @@ class KVCache(_BaseCache):
     def make_mask(self, *args, **kwargs):
         return create_attention_mask(*args, offset=self.offset, **kwargs)
 
+    @classmethod
+    def merge(_, caches):
+        return BatchKVCache.merge(caches)
+
+    def empty(self):
+        return self.keys is None
+
 
 class RotatingKVCache(_BaseCache):
+    step = 256
 
-    def __init__(self, max_size=None, keep=0, step=256):
+    def __init__(self, max_size, keep=0):
         self.keep = keep
         self.keys = None
         self.values = None
         self.offset = 0
         self.max_size = max_size
-        self.step = step
         self._idx = 0
 
     def _trim(self, trim_size, v, append=None):
@@ -388,10 +430,11 @@ class RotatingKVCache(_BaseCache):
             # preserve context
             self.keys = self._temporal_order(self.keys)
             self.values = self._temporal_order(self.values)
+            self._idx = self.keys.shape[2]
 
-            # The largest size is self.max_size + S to ensure
+            # The largest size is self.max_size + S - 1 to ensure
             # every token gets at least self.max_size context
-            trim_size = self._idx - self.max_size
+            trim_size = self._idx - self.max_size + 1
             self.keys = self._trim(trim_size, self.keys, keys)
             self.values = self._trim(trim_size, self.values, values)
         self.offset += keys.shape[2]
@@ -446,6 +489,9 @@ class RotatingKVCache(_BaseCache):
             return self._update_in_place(keys, values)
         return self._update_concat(keys, values)
 
+    def size(self):
+        return min(self.offset, self.max_size)
+
     @property
     def state(self):
         if self.offset < self.keys.shape[2]:
@@ -459,13 +505,11 @@ class RotatingKVCache(_BaseCache):
 
     @property
     def meta_state(self):
-        return tuple(
-            map(str, (self.keep, self.max_size, self.step, self.offset, self._idx))
-        )
+        return tuple(map(str, (self.keep, self.max_size, self.offset, self._idx)))
 
     @meta_state.setter
     def meta_state(self, v):
-        self.keep, self.max_size, self.step, self.offset, self._idx = map(
+        self.keep, self.max_size, self.offset, self._idx = map(
             int,
             v,
         )
@@ -487,7 +531,7 @@ class RotatingKVCache(_BaseCache):
     ):
         if N > 1:
             window_size = window_size or self.max_size
-            offset = min(self.max_size, self.offset)
+            offset = min(self.max_size - 1, self.offset)
             if offset + N > window_size or return_array:
                 return create_causal_mask(N, offset, window_size=window_size)
             else:
@@ -500,16 +544,27 @@ class RotatingKVCache(_BaseCache):
                 idx = self._idx
                 if idx >= self.max_size:
                     idx = 0
-                mask_size = min(self.max_size, self.offset)
+                if self.offset < self.max_size:
+                    mask_size = self.offset + 1
+                else:
+                    mask_size = self.max_size
                 mask = mx.arange(mask_size) >= (mask_size - window_size)
                 mask = mx.roll(mask, shift=idx + 1)
-                return mask[:, None]
+                return mask
+
+    @classmethod
+    def merge(_, caches):
+        return BatchRotatingKVCache.merge(caches)
+
+    def empty(self):
+        return self.keys is None
 
 
 class ArraysCache(_BaseCache):
     def __init__(self, size, left_padding: Optional[List[int]] = None):
         self.cache = [None] * size
-        self.left_padding = left_padding
+        self.left_padding = mx.array(left_padding) if left_padding else None
+        self.lengths = None
 
     def __setitem__(self, idx, value):
         self.cache[idx] = value
@@ -530,30 +585,73 @@ class ArraysCache(_BaseCache):
         In-place filter to keep just the given indices in the cache.
         """
         self.cache = [c[batch_indices] for c in self.cache]
-        self.left_padding = None
 
     def extend(self, other):
         """
         In-place extend this cache with the other cache.
         """
         self.cache = [mx.concatenate([c, o]) for c, o in zip(self.cache, other.cache)]
+
+    def extract(self, idx):
+        cache = ArraysCache(len(self.cache))
+        cache.cache = [c[idx : idx + 1] for c in self.cache]
+        return cache
+
+    def prepare(self, lengths=None, **kwargs):
+        self.lengths = mx.array(lengths)
+
+    def finalize(self):
+        self.lengths = None
         self.left_padding = None
 
+    def advance(self, N):
+        if self.lengths is not None:
+            self.lengths -= N
+        if self.left_padding is not None:
+            self.left_padding -= N
+
     def make_mask(self, N: int):
-        if self.cache[0] is None and self.left_padding is not None:
-            return mx.arange(N) >= self.left_padding[:, None]
+        if self.left_padding is not None:
+            pos = mx.arange(N)
+            return pos >= self.left_padding[:, None]
+        elif self.lengths is not None:
+            pos = mx.arange(N)
+            return pos < self.lengths[:, None]
         else:
             return None
 
+    @classmethod
+    def merge(cls, caches):
+        n_state = len(caches[0].cache)
+        B = len(caches)
+        cache = cls(n_state)
+        for e in range(n_state):
+            c_init = next(iter(c[e] for c in caches if c[e] is not None))
+            shape = list(c_init.shape)
+            shape[0] = B
+            cache[e] = mx.zeros(shape, c_init.dtype)
+            for i in range(B):
+                if caches[i][e] is None:
+                    continue
+                cache[e][i : i + 1] = caches[i][e]
+        return cache
+
+    def empty(self):
+        return self.cache[0] is None
+
 
 class MambaCache(ArraysCache):
     def __init__(self, left_padding: Optional[List[int]] = None):
         super().__init__(size=2, left_padding=left_padding)
 
 
-class ChunkedKVCache(KVCache):
-    def __init__(self, chunk_size=None):
-        super().__init__()
+class ChunkedKVCache(_BaseCache):
+    step = 256
+
+    def __init__(self, chunk_size):
+        self.keys = None
+        self.values = None
+        self.offset = 0
         self.chunk_size = chunk_size
         self.start_position = 0
 
@@ -589,6 +687,24 @@ class ChunkedKVCache(KVCache):
         self.values[..., prev:end, :] = values
         return self.keys[..., :end, :], self.values[..., :end, :]
 
+    @property
+    def state(self):
+        if self.offset == self.keys.shape[2]:
+            return self.keys, self.values
+        else:
+            return (
+                self.keys[..., : self.offset, :],
+                self.values[..., : self.offset, :],
+            )
+
+    @state.setter
+    def state(self, v):
+        self.keys, self.values = v
+        self.offset = self.keys.shape[2]
+
+    def is_trimmable(self):
+        return True
+
     def trim(self, n):
         n = min(self.offset - self.start_position, n)
         self.offset -= n
@@ -602,8 +718,11 @@ class ChunkedKVCache(KVCache):
     def meta_state(self, v):
         self.chunk_size, self.start_position = map(int, v)
 
+    def empty(self):
+        return self.keys is None
 
-class CacheList(KVCache):
+
+class CacheList(_BaseCache):
     def __init__(self, *caches):
         self.caches = caches
 
@@ -631,8 +750,59 @@ class CacheList(KVCache):
             c.state = v[start : start + l]
             start += l
 
+    def filter(self, batch_indices):
+        """
+        In-place filter to keep just the given indices in the cache.
+        """
+        for c in self.caches:
+            c.filter(batch_indices)
+
+    def extend(self, other):
+        """
+        In-place extend this cache with the other cache.
+        """
+        for c, o in zip(self.caches, other.caches):
+            c.extend(o)
+
+    @classmethod
+    def merge(cls, caches):
+        cache = cls()
+        cache.caches = tuple(
+            caches[0].caches[i].merge([c.caches[i] for c in caches])
+            for i in range(len(caches[0].caches))
+        )
+        return cache
+
+    def extract(self, idx):
+        return CacheList(*(c.extract(idx) for c in self.caches))
+
+    def prepare(self, **kwargs):
+        for c in self.caches:
+            c.prepare(**kwargs)
+
+    def finalize(self):
+        for c in self.caches:
+            c.finalize()
+
+    def size(self):
+        return max(c.size() for c in self.caches)
+
+    def empty(self):
+        return self.caches[0].empty()
+
+
+def dynamic_roll(x, shifts, axis):
+    n = x.shape[axis]
+    expand_shifts = (...,) + (None,) * (x.ndim - axis)
+    expand_indices = expand_shifts[:-1]
+    idx = (mx.arange(n)[expand_indices] - shifts[expand_shifts]) % n
+    rolled = mx.take_along_axis(x, idx, axis=axis)
+    return rolled
+
 
 class BatchKVCache(_BaseCache):
+    step = 256
+
     def __init__(self, left_padding: List[int]):
         """
         The BatchKV cache expects inputs to be left-padded.
@@ -657,7 +827,8 @@ class BatchKVCache(_BaseCache):
         self.left_padding = mx.array(left_padding)
         self.offset = mx.array([-l for l in left_padding])
         self._idx = 0
-        self.step = 256
+
+        self._right_padding = None
 
     def update_and_fetch(self, keys, values):
         prev = self._idx
@@ -684,6 +855,28 @@ class BatchKVCache(_BaseCache):
         self.values[..., prev : self._idx, :] = values
         return self.keys[..., : self._idx, :], self.values[..., : self._idx, :]
 
+    def prepare(self, *, left_padding=None, lengths=None, right_padding=None):
+        if left_padding is not None:
+            if self.keys is not None:
+                raise ValueError(
+                    "Left padding can only be added to an empty BatchKVCache"
+                )
+            left_padding = mx.array(left_padding)
+            self.left_padding += left_padding
+            self.offset -= left_padding
+
+        if right_padding is not None and max(right_padding) > 0:
+            self._right_padding = mx.array(right_padding)
+
+    def finalize(self):
+        if self._right_padding is not None:
+            padding = self._right_padding
+            self.keys = dynamic_roll(self.keys, padding[:, None], axis=2)
+            self.values = dynamic_roll(self.values, padding[:, None], axis=2)
+            self.offset -= padding
+            self.left_padding += padding
+            self._right_padding = None
+
     @property
     def state(self):
         k, v = self.keys, self.values
@@ -756,3 +949,356 @@ class BatchKVCache(_BaseCache):
             mx.concatenate, zip(*(pad(self), pad(other)))
         )
         self._idx = max_idx
+
+    def extract(self, idx):
+        cache = KVCache()
+        padding = self.left_padding[idx].item()
+        cache.keys = mx.contiguous(self.keys[idx : idx + 1, :, padding : self._idx])
+        cache.values = mx.contiguous(self.values[idx : idx + 1, :, padding : self._idx])
+        cache.offset = cache.keys.shape[2]
+        return cache
+
+    @classmethod
+    def merge(cls, caches):
+        lengths = [c.size() for c in caches]
+        max_length = max(lengths)
+        padding = [max_length - l for l in lengths]
+        B = len(caches)
+        H = max(c.keys.shape[1] for c in caches if c.keys is not None)
+        Dk = max(c.keys.shape[3] for c in caches if c.keys is not None)
+        Dv = max(c.values.shape[3] for c in caches if c.values is not None)
+        dt = next(iter(c.keys.dtype for c in caches if c.keys is not None))
+
+        keys = mx.zeros((B, H, max_length, Dk), dtype=dt)
+        values = mx.zeros((B, H, max_length, Dv), dtype=dt)
+        for i, (p, c) in enumerate(zip(padding, caches)):
+            if c.keys is None:
+                continue
+            keys[i : i + 1, :, p : p + c.offset] = c.keys[..., : c.offset, :]
+            values[i : i + 1, :, p : p + c.offset] = c.values[..., : c.offset, :]
+
+        cache = cls(padding)
+        cache.keys = keys
+        cache.values = values
+        cache.offset += keys.shape[2]
+        cache._idx = keys.shape[2]
+
+        return cache
+
+    def empty(self):
+        return self.keys is None
+
+
+class BatchRotatingKVCache(_BaseCache):
+    step = 256
+
+    def __init__(self, max_size, left_padding: List[int]):
+        self.keys = None
+        self.values = None
+
+        self.left_padding = mx.array(left_padding)
+        self.offset = mx.array([-l for l in left_padding])
+
+        self.max_size = max_size
+        self._idx = 0
+        self._offset = 0
+        self.rotated = False
+
+        # Lengths for right_padded inputs to make sure that padding tokens do
+        # not evict valid tokens.
+        self._lengths = None
+
+    def _trim(self, trim_size, v, append=None):
+        if trim_size > 0:
+            v = v[..., trim_size:, :]
+        if append is not None:
+            return mx.concatenate([v, append], axis=2)
+        return v
+
+    def _temporal_order(self):
+        """
+        Rearrange the cache into temporal order.
+        """
+        if self.rotated:
+            self.keys = mx.roll(self.keys, -self._idx, axis=2)
+            self.values = mx.roll(self.values, -self._idx, axis=2)
+            self._idx = self.keys.shape[2]
+            self.rotated = False
+
+    def _update_concat(self, keys, values):
+        if self.keys is None:
+            self.keys = keys
+            self.values = values
+        else:
+            # Put the keys/values in temporal order to
+            # preserve context
+            self._temporal_order()
+
+            # Slice off the end if needed
+            if self.keys.shape[2] > self._idx:
+                self.keys = self.keys[..., : self._idx, :]
+                self.values = self.values[..., : self._idx, :]
+
+            # Roll right sequences that are padded to make sure that we don't
+            # trim valid cache entries
+            if self._lengths is not None:
+                roll = mx.maximum(0, self.offset - self._lengths)
+                self.keys = dynamic_roll(self.keys, roll[:, None], axis=2)
+                self.values = dynamic_roll(self.values, roll[:, None], axis=2)
+                self.left_padding += roll
+                self.offset -= roll
+
+            # The largest size is self.max_size + S - 1 to ensure
+            # every token gets at least self.max_size context
+            trim_size = self._idx - self.max_size + 1
+            if trim_size > 0:
+                self.left_padding -= trim_size
+            self.keys = self._trim(trim_size, self.keys, keys)
+            self.values = self._trim(trim_size, self.values, values)
+        self.offset += keys.shape[2]
+        self._offset += keys.shape[2]
+        self._idx = self.keys.shape[2]
+        return self.keys, self.values
+
+    def _update_in_place(self, keys, values):
+        if self._lengths is not None:
+            raise RuntimeError(
+                "finalize() should be called before deocoding with BatchRotatingKVCache"
+            )
+
+        # May not have hit the max size yet, so potentially
+        # keep growing the cache
+        B, n_kv_heads, S, k_head_dim = keys.shape
+        prev = self._offset
+        if self.keys is None or (
+            prev >= self.keys.shape[2] and self.keys.shape[2] < self.max_size
+        ):
+            v_head_dim = values.shape[3]
+            new_size = min(self.step, self.max_size - prev)
+            k_shape = (B, n_kv_heads, new_size, k_head_dim)
+            v_shape = (B, n_kv_heads, new_size, v_head_dim)
+            new_k = mx.zeros(k_shape, keys.dtype)
+            new_v = mx.zeros(v_shape, values.dtype)
+            if self.keys is not None:
+                self.keys = mx.concatenate([self.keys, new_k], axis=2)
+                self.values = mx.concatenate([self.values, new_v], axis=2)
+            else:
+                self.keys, self.values = new_k, new_v
+            self._idx = prev
+
+        # Trim if needed
+        trim_size = self.keys.shape[2] - self.max_size
+        if trim_size > 0:
+            self.keys = self._trim(trim_size, self.keys)
+            self.values = self._trim(trim_size, self.values)
+            self._idx = self.max_size
+            self.left_padding -= trim_size
+
+        # Rotate
+        if self._idx == self.max_size:
+            self.rotated = True
+            self._idx = 0
+        if self.rotated:
+            self.left_padding -= S
+
+        # Assign
+        self.keys[..., self._idx : self._idx + S, :] = keys
+        self.values[..., self._idx : self._idx + S, :] = values
+        self._offset += S
+        self.offset += S
+        self._idx += S
+
+        # If the buffer is not full, slice off the end
+        if self._offset < self.max_size:
+            return (
+                self.keys[..., : self._offset, :],
+                self.values[..., : self._offset, :],
+            )
+        return self.keys, self.values
+
+    def update_and_fetch(self, keys, values):
+        if keys.shape[2] == 1:
+            return self._update_in_place(keys, values)
+        return self._update_concat(keys, values)
+
+    def prepare(self, *, left_padding=None, lengths=None, right_padding=None):
+        if left_padding is not None:
+            if self.keys is not None:
+                raise ValueError(
+                    "Left padding can only be added to an empty BatchRotatingKVCache"
+                )
+            left_padding = mx.array(left_padding)
+            self.left_padding += left_padding
+            self.offset -= left_padding
+
+        if right_padding is not None and max(right_padding) > 0:
+            self._lengths = mx.array(lengths) + self.offset
+
+    def finalize(self):
+        if self._lengths is not None:
+            roll = mx.maximum(0, self.offset - self._lengths)
+            self.keys = dynamic_roll(self.keys, roll[:, None], axis=2)
+            self.values = dynamic_roll(self.values, roll[:, None], axis=2)
+            self.left_padding += roll
+            self.offset -= roll
+            self._lengths = None
+
+    @property
+    def state(self):
+        k, v = self.keys, self.values
+        if self._offset < k.shape[2]:
+            k, v = k[..., : self._offset, :], v[..., : self._offset, :]
+        return k, v, self.offset, self.left_padding
+
+    @state.setter
+    def state(self, v):
+        self.keys, self.values, self.offset, self.left_padding = v
+
+    @property
+    def meta_state(self):
+        return tuple(map(str, (self.max_size, self._offset, self._idx, self.rotated)))
+
+    @meta_state.setter
+    def meta_state(self, v):
+        self.max_size, self._offset, self._idx = map(
+            int,
+            v[:3],
+        )
+        self.rotated = bool(v[3])
+
+    def is_trimmable(self):
+        return self._offset < self.max_size
+
+    def trim(self, n):
+        n = min(self._offset, n)
+        self._offset -= n
+        self._idx -= n
+        self.offset -= n
+        return n
+
+    def to_quantized(self, group_size: int = 64, bits: int = 4) -> QuantizedKVCache:
+        raise NotImplementedError("BatchRotatingKVCache Quantization NYI")
+
+    def make_mask(
+        self, N: int, window_size: Optional[int] = None, return_array: bool = False
+    ):
+        left_padding = self.left_padding
+        window_size = window_size or self.max_size
+        offset = min(self.max_size - 1, self._offset)
+        rinds = mx.arange(offset + N)
+        linds = mx.arange(offset, offset + N) if offset else rinds
+        linds = linds[:, None]
+        rinds = rinds[None]
+        mask = linds >= rinds
+        mask &= linds < rinds + window_size
+        if (trim_size := self._idx - self.max_size + int(N > 1)) > 0:
+            left_padding = left_padding - trim_size
+
+        rotated = N == 1 and (self.rotated or self._idx >= self.max_size)
+        if rotated:
+            left_padding = left_padding - 1
+
+        mask = mask & (rinds >= mx.expand_dims(left_padding, (1, 2, 3)))
+
+        if rotated:
+            idx = self._idx
+            if idx >= self.max_size:
+                idx = 0
+            mask = mx.roll(mask, shift=idx + 1, axis=-1)
+
+        return mask
+
+    def filter(self, batch_indices):
+        """
+        In-place filter to keep just the given indices in the cache.
+        """
+        self.keys = self.keys[batch_indices]
+        self.values = self.values[batch_indices]
+        self.offset = self.offset[batch_indices]
+        self.left_padding = self.left_padding[batch_indices]
+
+    def extend(self, other):
+        """
+        In-place extend this cache with the other cache.
+        """
+        if (self.rotated != other.rotated) or self._idx != other._idx:
+            self._temporal_order()
+            other._temporal_order()
+
+        max_idx = max(self._idx, other._idx)
+        max_size = max(self.keys.shape[2], other.keys.shape[2])
+
+        def pad(c):
+            left = max_idx - c._idx
+            right = max_size - c.keys.shape[2] - left
+            k, v = c.keys, c.values
+            if right < 0:
+                k = k[..., :right, :]
+                v = v[..., :right, :]
+                right = 0
+            if left != 0 or right != 0:
+                pad = [(0, 0), (0, 0), (left, right), (0, 0)]
+                k = mx.pad(k, pad)
+                v = mx.pad(v, pad)
+            left_padding = c.left_padding + left
+            return k, v, c.offset, left_padding
+
+        self.keys, self.values, self.offset, self.left_padding = map(
+            mx.concatenate, zip(*(pad(self), pad(other)))
+        )
+        self._idx = max_idx
+        self._offset = max(self._offset, other._offset)
+
+    def extract(self, idx):
+        cache = RotatingKVCache(self.max_size)
+        padding = self.left_padding[idx].item()
+        offset = self.offset[idx].item()
+        cache.keys = self.keys[idx : idx + 1]
+        cache.values = self.values[idx : idx + 1]
+        cache._idx = self._idx
+        if self.rotated:
+            cache.keys = mx.roll(cache.keys, -self._idx, axis=2)
+            cache.values = mx.roll(cache.values, -self._idx, axis=2)
+            cache._idx = self.max_size
+        cache.keys = mx.contiguous(cache.keys[:, :, padding : cache._idx])
+        cache.values = mx.contiguous(cache.values[:, :, padding : cache._idx])
+        cache.offset = offset
+        cache._idx = cache.keys.shape[2]
+        return cache
+
+    @classmethod
+    def merge(cls, caches):
+        if not all(c.max_size == caches[0].max_size for c in caches):
+            raise ValueError(
+                "BatchRotatingKVCache can only merge caches with the same maximum size"
+            )
+
+        offsets = [c.offset for c in caches]
+        lengths = [c.size() for c in caches]
+        max_length = max(lengths)
+        padding = [max_length - l for l in lengths]
+        B = len(caches)
+        H = max(c.keys.shape[1] for c in caches if c.keys is not None)
+        Dk = max(c.keys.shape[3] for c in caches if c.keys is not None)
+        Dv = max(c.values.shape[3] for c in caches if c.values is not None)
+        dt = next(iter(c.keys.dtype for c in caches if c.keys is not None))
+
+        keys = mx.zeros((B, H, max_length, Dk), dtype=dt)
+        values = mx.zeros((B, H, max_length, Dv), dtype=dt)
+        for i, (p, c) in enumerate(zip(padding, caches)):
+            if c.keys is None:
+                continue
+            keys[i : i + 1, :, p : p + c._idx] = c._temporal_order(c.keys)
+            values[i : i + 1, :, p : p + c._idx] = c._temporal_order(c.values)
+
+        cache = cls(caches[0].max_size, padding)
+        cache.keys = keys
+        cache.values = values
+        cache.offset = mx.array(offsets)
+        cache._idx = keys.shape[2]
+        cache._offset = keys.shape[2]
+
+        return cache
+
+    def empty(self):
+        return self.keys is None

mlx_lm/models/cohere.py

@@ -6,6 +6,7 @@ from typing import Any, Optional
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
@@ -109,7 +110,7 @@ class MLP(nn.Module):
         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))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class TransformerBlock(nn.Module):

mlx_lm/models/cohere2.py

@@ -6,6 +6,7 @@ from typing import Optional, Tuple
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .cache import KVCache, RotatingKVCache
 
@@ -106,7 +107,7 @@ class MLP(nn.Module):
         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))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class TransformerBlock(nn.Module):

mlx_lm/models/dbrx.py

@@ -7,6 +7,7 @@ import mlx.core as mx
 import mlx.nn as nn
 import numpy as np
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
@@ -107,7 +108,7 @@ class MLP(nn.Module):
         self.w2 = nn.Linear(ffn_dim, d_model, bias=False)
 
     def __call__(self, x: mx.array) -> mx.array:
-        current_hidden_states = nn.silu(self.w1(x)) * self.v1(x)
+        current_hidden_states = swiglu(self.w1(x), self.v1(x))
         current_hidden_states = self.w2(current_hidden_states)
         return current_hidden_states
 

mlx_lm/models/deepseek.py

@@ -4,6 +4,7 @@ from typing import Any, Dict, Optional
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .switch_layers import SwitchGLU
 
@@ -120,7 +121,7 @@ class DeepseekMLP(nn.Module):
         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))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class MoEGate(nn.Module):

mlx_lm/models/deepseek_v2.py

@@ -6,8 +6,11 @@ from typing import Any, Dict, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
+from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .pipeline import PipelineMixin
 from .switch_layers import SwitchGLU
 
 
@@ -258,7 +261,7 @@ class DeepseekV2MLP(nn.Module):
         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))
+        down_proj = self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
         return down_proj
 
 
@@ -314,13 +317,21 @@ class DeepseekV2MoE(nn.Module):
                 config=config, intermediate_size=intermediate_size
             )
 
+        self.sharding_group = None
+
     def __call__(self, x):
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(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)
 
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+
         return y
 
 
@@ -355,7 +366,7 @@ class DeepseekV2DecoderLayer(nn.Module):
         return out
 
 
-class DeepseekV2Model(nn.Module):
+class DeepseekV2Model(PipelineMixin, nn.Module):
     def __init__(self, config: ModelArgs):
         super().__init__()
         self.vocab_size = config.vocab_size
@@ -364,32 +375,8 @@ class DeepseekV2Model(nn.Module):
             DeepseekV2DecoderLayer(config, idx)
             for idx in range(config.num_hidden_layers)
         ]
-        self.start_idx = 0
-        self.end_idx = len(self.layers)
-        self.num_layers = self.end_idx
-
         self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
-        self.pipeline_rank = 0
-        self.pipeline_size = 1
-
-    def pipeline(self, group):
-        # Split layers in reverse so rank=0 gets the last layers and
-        # rank=pipeline_size-1 gets the first
-        self.pipeline_rank = group.rank()
-        self.pipeline_size = group.size()
-        layers_per_rank = len(self.layers) // self.pipeline_size
-        extra = len(self.layers) - layers_per_rank * self.pipeline_size
-        if self.pipeline_rank < extra:
-            layers_per_rank += 1
-
-        self.start_idx = (self.pipeline_size - self.pipeline_rank - 1) * layers_per_rank
-        self.end_idx = self.start_idx + layers_per_rank
-        self.num_layers = layers_per_rank
-        self.layers = self.layers[: self.end_idx]
-        self.layers[: self.start_idx] = [None] * self.start_idx
-        self.num_layers = len(self.layers) - self.start_idx
-
     def __call__(
         self,
         x: mx.array,
@@ -401,22 +388,25 @@ class DeepseekV2Model(nn.Module):
         pipeline_size = self.pipeline_size
 
         if cache is None:
-            cache = [None] * self.num_layers
+            cache = [None] * len(self.pipeline_layers)
         mask = create_attention_mask(h, cache[0])
 
         # Receive from the previous process in the pipeline
         if pipeline_rank < pipeline_size - 1:
             h = mx.distributed.recv_like(h, (pipeline_rank + 1))
 
-        for i in range(self.num_layers):
-            h = self.layers[self.start_idx + i](h, mask, cache[i])
+        for l, c in zip(self.pipeline_layers, cache):
+            h = l(h, mask, cache=c)
 
         # Send to the next process in the pipeline
         if pipeline_rank != 0:
             h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
+            if cache[-1] is not None:
+                cache[-1].keys = mx.depends(cache[-1].keys, h)
 
         # Broadcast h while keeping it in the graph
-        h = mx.distributed.all_gather(h)[: h.shape[0]]
+        if pipeline_size > 1:
+            h = mx.distributed.all_gather(h)[: h.shape[0]]
 
         return self.norm(h)
 
@@ -450,6 +440,62 @@ class Model(nn.Module):
                         weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
         return weights
 
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        for layer in self.model.layers:
+            # Shard the self attention
+            if layer.self_attn.q_lora_rank is None:
+                layer.self_attn.q_proj = shard_linear(
+                    layer.self_attn.q_proj, "all-to-sharded", group=group
+                )
+            else:
+                layer.self_attn.q_b_proj = shard_linear(
+                    layer.self_attn.q_b_proj, "all-to-sharded", group=group
+                )
+            layer.self_attn.kv_b_proj = shard_linear(
+                layer.self_attn.kv_b_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.o_proj = shard_linear(
+                layer.self_attn.o_proj, "sharded-to-all", group=group
+            )
+            layer.self_attn.num_heads //= N
+
+            # Shard the MLP
+            if isinstance(layer.mlp, DeepseekV2MLP):
+                layer.mlp.gate_proj = shard_linear(
+                    layer.mlp.gate_proj, "all-to-sharded", group=group
+                )
+                layer.mlp.down_proj = shard_linear(
+                    layer.mlp.down_proj, "sharded-to-all", group=group
+                )
+                layer.mlp.up_proj = shard_linear(
+                    layer.mlp.up_proj, "all-to-sharded", group=group
+                )
+
+            # Shard the MoE. Shard in place since the MoE should be responsible
+            # for aggregating the results.
+            else:
+                layer.mlp.sharding_group = group
+                shard_inplace(
+                    layer.mlp.shared_experts.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_experts.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_experts.up_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.up_proj, "all-to-sharded", group=group
+                )
+
     @property
     def layers(self):
-        return self.model.layers[self.model.start_idx : self.model.end_idx]
+        return self.model.pipeline_layers

mlx_lm/models/deepseek_v3.py

@@ -7,8 +7,12 @@ from typing import Any, Dict, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
+from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .pipeline import PipelineMixin
+from .rope_utils import initialize_rope
 from .switch_layers import SwitchGLU
 
 
@@ -45,85 +49,6 @@ class ModelArgs(BaseModelArgs):
     attention_bias: bool = False
 
 
-def yarn_find_correction_dim(
-    num_rotations, dim, base=10000, max_position_embeddings=2048
-):
-    return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / (
-        2 * math.log(base)
-    )
-
-
-def yarn_find_correction_range(
-    low_rot, high_rot, dim, base=10000, max_position_embeddings=2048
-):
-    low = math.floor(
-        yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings)
-    )
-    high = math.ceil(
-        yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings)
-    )
-    return max(low, 0), min(high, dim - 1)
-
-
-def yarn_get_mscale(scale=1, mscale=1):
-    if scale <= 1:
-        return 1.0
-    return 0.1 * mscale * math.log(scale) + 1.0
-
-
-def yarn_linear_ramp_mask(min_val, max_val, dim):
-    if min_val == max_val:
-        max_val += 0.001  # Prevent singularity
-
-    linear_func = (mx.arange(dim, dtype=mx.float32) - min_val) / (max_val - min_val)
-    return mx.clip(linear_func, 0, 1)
-
-
-class DeepseekV3YarnRotaryEmbedding(nn.Module):
-    def __init__(
-        self,
-        dim,
-        max_position_embeddings=2048,
-        base=10000,
-        scaling_factor=1.0,
-        original_max_position_embeddings=4096,
-        beta_fast=32,
-        beta_slow=1,
-        mscale=1,
-        mscale_all_dim=0,
-    ):
-        super().__init__()
-        self.mscale = yarn_get_mscale(scaling_factor, mscale) / yarn_get_mscale(
-            scaling_factor, mscale_all_dim
-        )
-        freq_extra = base ** (mx.arange(0, dim, 2, dtype=mx.float32) / dim)
-        freq_inter = scaling_factor * freq_extra
-        low, high = yarn_find_correction_range(
-            beta_fast,
-            beta_slow,
-            dim,
-            base,
-            original_max_position_embeddings,
-        )
-        freq_mask = 1.0 - yarn_linear_ramp_mask(low, high, dim // 2)
-        self._freqs = (freq_inter * freq_extra) / (
-            freq_inter * freq_mask + freq_extra * (1 - freq_mask)
-        )
-
-    def __call__(self, x, offset=0):
-        if self.mscale != 1.0:
-            x = self.mscale * x
-        return mx.fast.rope(
-            x,
-            x.shape[-1],
-            traditional=True,
-            base=None,
-            scale=1.0,
-            offset=offset,
-            freqs=self._freqs,
-        )
-
-
 class DeepseekV3Attention(nn.Module):
     def __init__(self, config: ModelArgs):
         super().__init__()
@@ -175,35 +100,19 @@ class DeepseekV3Attention(nn.Module):
 
         if self.config.rope_scaling is not None:
             mscale_all_dim = self.config.rope_scaling.get("mscale_all_dim", 0)
-            scaling_factor = self.config.rope_scaling["factor"]
             if mscale_all_dim:
-                mscale = yarn_get_mscale(scaling_factor, mscale_all_dim)
-                self.scale = self.scale * mscale * mscale
+                scaling_factor = self.config.rope_scaling["factor"]
+                if scaling_factor > 1:
+                    s = 0.1 * mscale_all_dim * math.log(scaling_factor) + 1.0
+                    self.scale = self.scale * s * s
 
-            rope_kwargs = {
-                key: self.config.rope_scaling[key]
-                for key in [
-                    "original_max_position_embeddings",
-                    "beta_fast",
-                    "beta_slow",
-                    "mscale",
-                    "mscale_all_dim",
-                ]
-                if key in self.config.rope_scaling
-            }
-            self.rope = DeepseekV3YarnRotaryEmbedding(
-                dim=self.qk_rope_head_dim,
-                max_position_embeddings=self.max_position_embeddings,
-                scaling_factor=scaling_factor,
-                base=self.rope_theta,
-                **rope_kwargs,
-            )
-        else:
-            self.rope = nn.RoPE(
-                dims=self.qk_rope_head_dim,
-                base=self.rope_theta,
-                traditional=True,
-            )
+        self.rope = initialize_rope(
+            dims=self.qk_rope_head_dim,
+            base=self.rope_theta,
+            traditional=True,
+            max_position_embeddings=self.max_position_embeddings,
+            scaling_config=self.config.rope_scaling,
+        )
 
     def __call__(
         self,
@@ -266,7 +175,7 @@ class DeepseekV3MLP(nn.Module):
         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))
+        down_proj = self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
         return down_proj
 
 
@@ -349,13 +258,21 @@ class DeepseekV3MoE(nn.Module):
                 config=config, intermediate_size=intermediate_size
             )
 
+        self.sharding_group = None
+
     def __call__(self, x):
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(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)
 
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+
         return y
 
 
@@ -389,7 +306,7 @@ class DeepseekV3DecoderLayer(nn.Module):
         return h + r
 
 
-class DeepseekV3Model(nn.Module):
+class DeepseekV3Model(PipelineMixin, nn.Module):
     def __init__(self, config: ModelArgs):
         super().__init__()
         self.vocab_size = config.vocab_size
@@ -398,28 +315,7 @@ class DeepseekV3Model(nn.Module):
             DeepseekV3DecoderLayer(config, idx)
             for idx in range(config.num_hidden_layers)
         ]
-        self.start_idx = 0
-        self.end_idx = len(self.layers)
-        self.num_layers = self.end_idx
-
         self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
-        self.pipeline_rank = 0
-        self.pipeline_size = 1
-
-    def pipeline(self, group):
-        # Split layers in reverse so rank=0 gets the last layers and
-        # rank=pipeline_size-1 gets the first
-        self.pipeline_rank = group.rank()
-        self.pipeline_size = group.size()
-        layers_per_rank = len(self.layers) // self.pipeline_size
-        extra = len(self.layers) - layers_per_rank * self.pipeline_size
-        if self.pipeline_rank < extra:
-            layers_per_rank += 1
-        self.start_idx = (self.pipeline_size - self.pipeline_rank - 1) * layers_per_rank
-        self.end_idx = self.start_idx + layers_per_rank
-        self.layers = self.layers[: self.end_idx]
-        self.layers[: self.start_idx] = [None] * self.start_idx
-        self.num_layers = len(self.layers) - self.start_idx
 
     def __call__(
         self,
@@ -432,23 +328,25 @@ class DeepseekV3Model(nn.Module):
         pipeline_size = self.pipeline_size
 
         if cache is None:
-            cache = [None] * self.num_layers
+            cache = [None] * len(self.pipeline_layers)
         mask = create_attention_mask(h, cache[0])
 
         # Receive from the previous process in the pipeline
-
         if pipeline_rank < pipeline_size - 1:
             h = mx.distributed.recv_like(h, (pipeline_rank + 1))
 
-        for i in range(self.num_layers):
-            h = self.layers[self.start_idx + i](h, mask, cache[i])
+        for l, c in zip(self.pipeline_layers, cache):
+            h = l(h, mask, cache=c)
 
         # Send to the next process in the pipeline
         if pipeline_rank != 0:
             h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
+            if cache[-1] is not None:
+                cache[-1].keys = mx.depends(cache[-1].keys, h)
 
         # Broadcast h while keeping it in the graph
-        h = mx.distributed.all_gather(h)[: h.shape[0]]
+        if pipeline_size > 1:
+            h = mx.distributed.all_gather(h)[: h.shape[0]]
 
         return self.norm(h)
 
@@ -471,7 +369,8 @@ class Model(nn.Module):
 
     def sanitize(self, weights):
         def dequant(weight, scale_inv):
-            dtype = weight.dtype
+            dtype = mx.bfloat16
+            weight = mx.from_fp8(weight, dtype=mx.bfloat16)
             bs = 128  # block size
             m, n = weight.shape
             pad_bottom = (-m) % bs
@@ -485,7 +384,22 @@ class Model(nn.Module):
             )
             return weight[:m, :n].astype(dtype)
 
-        # Dequantize
+        # Remap for int4
+        new_weights = {}
+        for k, v in weights.items():
+            if k.endswith("weight_shape"):
+                base = k.replace("weight_shape", "")
+                new_weights[base + "weight"] = weights[base + "weight_packed"].view(
+                    mx.uint32
+                )
+                s = weights[base + "weight_scale"]
+                new_weights[base + "scales"] = s
+                new_weights[base + "biases"] = -8 * s
+            elif not (k.endswith("weight_scale") or k.endswith("weight_packed")):
+                new_weights[k] = v
+        weights = new_weights
+
+        # Dequantize fp8
         new_weights = {}
         for k, v in weights.items():
             if "weight_scale_inv" in k:
@@ -517,9 +431,65 @@ class Model(nn.Module):
             if not k.startswith("model.layers.61") and "rotary_emb.inv_freq" not in k
         }
 
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        for layer in self.model.layers:
+            # Shard the self attention
+            if layer.self_attn.q_lora_rank is None:
+                layer.self_attn.q_proj = shard_linear(
+                    layer.self_attn.q_proj, "all-to-sharded", group=group
+                )
+            else:
+                layer.self_attn.q_b_proj = shard_linear(
+                    layer.self_attn.q_b_proj, "all-to-sharded", group=group
+                )
+            layer.self_attn.kv_b_proj = shard_linear(
+                layer.self_attn.kv_b_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.o_proj = shard_linear(
+                layer.self_attn.o_proj, "sharded-to-all", group=group
+            )
+            layer.self_attn.num_heads //= N
+
+            # Shard the MLP
+            if isinstance(layer.mlp, DeepseekV3MLP):
+                layer.mlp.gate_proj = shard_linear(
+                    layer.mlp.gate_proj, "all-to-sharded", group=group
+                )
+                layer.mlp.down_proj = shard_linear(
+                    layer.mlp.down_proj, "sharded-to-all", group=group
+                )
+                layer.mlp.up_proj = shard_linear(
+                    layer.mlp.up_proj, "all-to-sharded", group=group
+                )
+
+            # Shard the MoE. Shard in place since the MoE should be responsible
+            # for aggregating the results.
+            else:
+                layer.mlp.sharding_group = group
+                shard_inplace(
+                    layer.mlp.shared_experts.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_experts.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_experts.up_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.up_proj, "all-to-sharded", group=group
+                )
+
     @property
     def layers(self):
-        return self.model.layers[self.model.start_idx : self.model.end_idx]
+        return self.model.pipeline_layers
 
     @property
     def cast_predicate(self):

mlx_lm/models/deepseek_v32.py

@@ -0,0 +1,582 @@
+# 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 mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
+
+from .activations import swiglu
+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
+        # Ensure the indexer cache is evaluated even if the topk_indices are unused
+        # to keep the graph from getting too large
+        if cache is not None and cache[0] is not None:
+            cache[0].keys = mx.depends(cache[0].keys, (cache[1].keys, cache[1].values))
+
+        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(swiglu(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
+            )
+
+        self.sharding_group = None
+
+    def __call__(self, x):
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(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)
+
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+
+        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][0].keys = mx.depends(cache[-1][0].keys, h)
+
+        # Broadcast h while keeping it in the graph
+        if pipeline_size > 1:
+            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 = mx.bfloat16
+            weight = mx.from_fp8(weight, dtype=mx.bfloat16)
+            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
+        }
+
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        for layer in self.model.layers:
+            layer.self_attn.q_b_proj = shard_linear(
+                layer.self_attn.q_b_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.kv_b_proj = shard_linear(
+                layer.self_attn.kv_b_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.o_proj = shard_linear(
+                layer.self_attn.o_proj, "sharded-to-all", group=group
+            )
+            layer.self_attn.num_heads //= N
+
+            # Shard the MLP
+            if isinstance(layer.mlp, DeepseekV32MLP):
+                layer.mlp.gate_proj = shard_linear(
+                    layer.mlp.gate_proj, "all-to-sharded", group=group
+                )
+                layer.mlp.down_proj = shard_linear(
+                    layer.mlp.down_proj, "sharded-to-all", group=group
+                )
+                layer.mlp.up_proj = shard_linear(
+                    layer.mlp.up_proj, "all-to-sharded", group=group
+                )
+
+            # Shard the MoE. Shard in place since the MoE should be responsible
+            # for aggregating the results.
+            else:
+                layer.mlp.sharding_group = group = group
+                shard_inplace(
+                    layer.mlp.shared_experts.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_experts.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_experts.up_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.up_proj, "all-to-sharded", group=group
+                )
+
+    @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

@@ -7,6 +7,7 @@ from typing import Any, Dict, Optional, Union
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .rope_utils import initialize_rope
 from .switch_layers import SwitchGLU
@@ -180,7 +181,7 @@ class Dots1MLP(nn.Module):
         )
 
     def __call__(self, x) -> mx.array:
-        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class Dots1MoE(nn.Module):

mlx_lm/models/ernie4_5.py

@@ -6,6 +6,7 @@ from typing import Any, Optional
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .rope_utils import initialize_rope
 
@@ -87,7 +88,7 @@ class MLP(nn.Module):
         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))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class DecoderLayer(nn.Module):

mlx_lm/models/ernie4_5_moe.py

@@ -6,6 +6,7 @@ from typing import Any, Optional
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .rope_utils import initialize_rope
 from .switch_layers import SwitchGLU
@@ -98,7 +99,7 @@ class Ernie4_5_MLP(nn.Module):
         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))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class Ernie4_5_MoeMLP(nn.Module):

mlx_lm/models/exaone.py

@@ -6,6 +6,7 @@ from typing import Any, Dict, Optional, Union
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .rope_utils import initialize_rope
 
@@ -91,7 +92,7 @@ class MLP(nn.Module):
         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))
+        return self.c_proj(swiglu(self.c_fc_0(x), self.c_fc_1(x)))
 
 
 class TransformerBlock(nn.Module):

mlx_lm/models/exaone4.py

@@ -6,6 +6,7 @@ from typing import Any, Dict, Optional, Union
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .cache import KVCache, RotatingKVCache
 from .rope_utils import initialize_rope
@@ -102,7 +103,7 @@ class MLP(nn.Module):
         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))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class TransformerBlock(nn.Module):

mlx_lm/models/exaone_moe.py

@@ -0,0 +1,439 @@
+# Copyright © 2026 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, List, Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
+
+from .activations import swiglu
+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
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    moe_intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    head_dim: int
+    num_experts: int
+    num_experts_per_tok: int
+    num_shared_experts: int
+    rms_norm_eps: float
+    max_position_embeddings: int
+    sliding_window: int
+    layer_types: List[str]
+    is_moe_layer: List[bool]
+    n_group: int = 1
+    topk_group: int = 1
+    routed_scaling_factor: float = 2.5
+    norm_topk_prob: bool = True
+    scoring_func: str = "sigmoid"
+    topk_method: str = "noaux_tc"
+    rope_theta: float = 1000000.0
+    rope_scaling: Optional[dict] = None
+    rope_parameters: Optional[dict] = None
+    tie_word_embeddings: bool = False
+
+    def __post_init__(self):
+        if self.rope_parameters is not None and "rope_theta" in self.rope_parameters:
+            self.rope_theta = self.rope_parameters["rope_theta"]
+
+
+@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 + 1e-20)
+
+    scores = scores * routed_scaling_factor
+    return inds, scores
+
+
+class MoEGate(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.num_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,))
+        assert args.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 MLP(nn.Module):
+    def __init__(self, args: ModelArgs, intermediate_size: Optional[int] = None):
+        super().__init__()
+        hidden_size = args.hidden_size
+        intermediate_size = intermediate_size or args.intermediate_size
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+
+    def __call__(self, x):
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
+
+
+class MoE(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.switch_mlp = SwitchGLU(
+            args.hidden_size,
+            args.moe_intermediate_size,
+            args.num_experts,
+        )
+
+        self.gate = MoEGate(args)
+
+        self.shared_experts = (
+            MLP(
+                args,
+                intermediate_size=args.moe_intermediate_size * args.num_shared_experts,
+            )
+            if args.num_shared_experts is not None and args.num_shared_experts > 0
+            else None
+        )
+
+        self.sharding_group = None
+
+    def __call__(self, x):
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(x)
+
+        inds, scores = self.gate(x)
+        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)
+
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+
+        return y
+
+
+class Attention(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.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(
+            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.is_sliding_window = args.layer_types[layer_idx] == "sliding_attention"
+        self.apply_rope_all_layers = "sliding_attention" not in args.layer_types
+        self.use_rope = self.is_sliding_window or self.apply_rope_all_layers
+
+        if self.use_rope:
+            self.rope = initialize_rope(
+                self.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 DecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+
+        self.self_attn = Attention(args, layer_idx)
+        self.mlp = MoE(args) if args.is_moe_layer[layer_idx] else MLP(args)
+        self.is_sliding_window = self.self_attn.is_sliding_window
+
+        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 ExaoneMoEModel(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 = [DecoderLayer(args, idx) for idx in range(args.num_hidden_layers)]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+        self.swa_idx = None
+        self.ga_idx = None
+        for i, layer in enumerate(self.layers):
+            if layer.is_sliding_window and self.swa_idx is None:
+                self.swa_idx = i
+            if not layer.is_sliding_window and self.ga_idx is None:
+                self.ga_idx = i
+
+        self.window_size = args.sliding_window
+
+    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)
+
+        global_mask = create_attention_mask(
+            h, cache[self.ga_idx] if self.ga_idx is not None else cache[0]
+        )
+        swa_mask = create_attention_mask(
+            h,
+            cache[self.swa_idx] if self.swa_idx is not None else cache[0],
+            window_size=self.window_size,
+        )
+
+        for layer, c in zip(self.layers, cache):
+            mask = swa_mask if layer.is_sliding_window 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 = ExaoneMoEModel(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,
+    ):
+        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):
+        new_weights = {k: v for k, v in weights.items() if not k.startswith("mtp.")}
+        weights = new_weights
+
+        for l in range(self.args.num_hidden_layers):
+            if not self.args.is_moe_layer[l]:
+                continue
+
+            prefix = f"model.layers.{l}"
+
+            bias_key = f"{prefix}.mlp.e_score_correction_bias"
+            if bias_key in weights:
+                weights[f"{prefix}.mlp.gate.e_score_correction_bias"] = weights.pop(
+                    bias_key
+                )
+
+            for m in ["gate_proj", "down_proj", "up_proj"]:
+                for k in ["weight", "scales", "biases"]:
+                    first_key = f"{prefix}.mlp.experts.0.{m}.{k}"
+                    last_key = (
+                        f"{prefix}.mlp.experts.{self.args.num_experts - 1}.{m}.{k}"
+                    )
+                    if first_key in weights and last_key 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 self.args.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
+
+        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
+
+    def make_cache(self):
+        caches = []
+        for layer in self.layers:
+            if layer.is_sliding_window:
+                caches.append(
+                    RotatingKVCache(max_size=self.args.sliding_window, keep=0)
+                )
+            else:
+                caches.append(KVCache())
+        return caches
+
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+
+        for layer in self.model.layers:
+            layer.self_attn.q_proj = shard_linear(
+                layer.self_attn.q_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.k_proj = shard_linear(
+                layer.self_attn.k_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.v_proj = shard_linear(
+                layer.self_attn.v_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.o_proj = shard_linear(
+                layer.self_attn.o_proj, "sharded-to-all", group=group
+            )
+            layer.self_attn.n_heads //= N
+            layer.self_attn.n_kv_heads //= N
+
+            if isinstance(layer.mlp, MLP):
+                layer.mlp.gate_proj = shard_linear(
+                    layer.mlp.gate_proj, "all-to-sharded", group=group
+                )
+                layer.mlp.down_proj = shard_linear(
+                    layer.mlp.down_proj, "sharded-to-all", group=group
+                )
+                layer.mlp.up_proj = shard_linear(
+                    layer.mlp.up_proj, "all-to-sharded", group=group
+                )
+            else:
+                layer.mlp.sharding_group = group
+                if layer.mlp.shared_experts is not None:
+                    shard_inplace(
+                        layer.mlp.shared_experts.gate_proj,
+                        "all-to-sharded",
+                        group=group,
+                    )
+                    shard_inplace(
+                        layer.mlp.shared_experts.down_proj,
+                        "sharded-to-all",
+                        group=group,
+                    )
+                    shard_inplace(
+                        layer.mlp.shared_experts.up_proj, "all-to-sharded", group=group
+                    )
+                shard_inplace(
+                    layer.mlp.switch_mlp.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.up_proj, "all-to-sharded", group=group
+                )

mlx_lm/models/falcon_h1.py

@@ -0,0 +1,504 @@
+# 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 .activations import swiglu
+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
+    tie_word_embeddings: bool = True
+
+
+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 = swiglu(gate, hidden_states)
+
+        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 = swiglu(gate, hidden_states)
+        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 _conv(
+        self,
+        conv_input: mx.array,
+        cache: Optional[MambaCache],
+        mask: Optional[mx.array],
+    ) -> mx.array:
+        if mask is not None:
+            conv_input = mx.where(mask[..., None], conv_input, 0)
+
+        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)
+            n_keep = self.conv_kernel_size - 1
+            if cache.lengths is not None:
+                t = padded_input.shape[1]
+                ends = mx.clip(cache.lengths, 0, t - n_keep)
+                positions = (ends[:, None] + mx.arange(n_keep))[..., None]
+                cache[0] = mx.take_along_axis(padded_input, positions, axis=1)
+            else:
+                cache[0] = padded_input[:, -n_keep:, :]
+        else:
+            padded_input = mx.pad(
+                conv_input, [(0, 0), (self.conv_kernel_size - 1, 0), (0, 0)]
+            )
+
+        conv_output = self.conv1d(padded_input)
+        return nn.silu(conv_output)
+
+    def _ssm(
+        self,
+        hidden_states: mx.array,
+        B: mx.array,
+        C: mx.array,
+        dt: mx.array,
+        cache: Optional[MambaCache],
+        mask: Optional[mx.array],
+    ) -> 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)
+        if cache:
+            state = cache[1]
+            lengths = cache.lengths
+        else:
+            state, lengths = None, None
+        y, state = ssm_update(
+            hidden_states,
+            self.A_log,
+            B,
+            C,
+            self.D,
+            dt,
+            self.dt_bias,
+            state,
+            self.time_step_limit,
+            mask,
+            lengths,
+        )
+        if cache:
+            cache[1] = state
+        return y.reshape(batch_size, seq_len, self.intermediate_size)
+
+    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,
+        )
+
+        conv_output = self._conv(conv_input, cache, mask)
+
+        hidden_states, B, C = mx.split(
+            conv_output,
+            [
+                self.intermediate_size,
+                self.intermediate_size + self.n_groups * self.ssm_state_size,
+            ],
+            axis=-1,
+        )
+
+        y = self._ssm(hidden_states, B, C, dt, cache, mask=mask)
+        if cache:
+            cache.advance(y.shape[1])
+
+        if self.mamba_rms_norm:
+            y = self.norm(y, gate)
+        else:
+            y = swiglu(gate, y)
+
+        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 = swiglu(self.gate_proj(x), self.up_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)
+        if not args.tie_word_embeddings:
+            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)
+        if self.args.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(hidden_states)
+            return out * (self.args.lm_head_multiplier / self.args.embedding_multiplier)
+        else:
+            return self.lm_head(hidden_states)
+
+    def sanitize(self, weights):
+        # Check if needs sanitization
+        c1d = weights["model.layers.0.mamba.conv1d.weight"]
+        if c1d.shape[-1] <= c1d.shape[1]:
+            return weights
+
+        sanitized_weights = {}
+        args = self.args
+
+        for name, param in weights.items():
+            # Fold-in multipliers
+            if name.endswith("embed_tokens.weight"):
+                param *= args.embedding_multiplier
+            elif name.endswith("lm_head.weight"):
+                param *= args.lm_head_multiplier
+            elif name.endswith("q_proj.weight") or name.endswith("k_proj.weight"):
+                param *= args.attention_in_multiplier
+            elif name.endswith("key_proj.weight"):
+                param *= args.attention_in_multiplier * args.key_multiplier
+            elif name.endswith("o_proj.weight"):
+                param *= args.attention_out_multiplier
+            elif name.endswith("out_proj.weight"):
+                param *= args.ssm_out_multiplier
+            elif name.endswith("gate_proj.weight"):
+                param *= args.mlp_multipliers[0]
+            elif name.endswith("down_proj.weight"):
+                param *= args.mlp_multipliers[1]
+            elif name.endswith("in_proj.weight"):
+                param *= (
+                    args.ssm_in_multiplier
+                    * self.model._mup_vector.astype(param.dtype)[:, None]
+                )
+            elif "conv1d.weight" in name:
+                param = param.transpose(0, 2, 1)
+            sanitized_weights[name] = param
+        return sanitized_weights
+
+    def make_cache(self):
+        return [
+            CacheList(MambaCache(), KVCache())
+            for _ in range(self.args.num_hidden_layers)
+        ]
+
+    @property
+    def layers(self):
+        return self.model.layers

mlx_lm/models/gated_delta.py

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

mlx_lm/models/gemma3_text.py

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

mlx_lm/models/glm.py

@@ -6,6 +6,7 @@ from typing import Any, Dict, Optional, Union
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .rope_utils import initialize_rope
 
@@ -102,7 +103,7 @@ class GLMMLP(nn.Module):
     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)
+        return self.down_proj(swiglu(gate, x))
 
 
 class GLMBlock(nn.Module):

mlx_lm/models/glm4.py

@@ -6,6 +6,7 @@ from typing import Any, Optional
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
@@ -38,7 +39,7 @@ class Glm4MLP(nn.Module):
     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)
+        return self.down_proj(swiglu(gate, up_states))
 
 
 class Glm4Attention(nn.Module):

mlx_lm/models/glm4_moe.py

@@ -7,8 +7,11 @@ from typing import Any, Dict, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
+from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .pipeline import PipelineMixin
 from .switch_layers import SwitchGLU
 
 
@@ -121,7 +124,7 @@ class MLP(nn.Module):
         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))
+        down_proj = self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
         return down_proj
 
 
@@ -204,13 +207,21 @@ class MoE(nn.Module):
                 config=config, intermediate_size=intermediate_size
             )
 
+        self.sharding_group = None
+
     def __call__(self, x):
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(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)
 
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+
         return y
 
 
@@ -243,7 +254,7 @@ class DecoderLayer(nn.Module):
         return h + r
 
 
-class LanguageModel(nn.Module):
+class LanguageModel(PipelineMixin, nn.Module):
     def __init__(self, config: ModelArgs):
         super().__init__()
         self.vocab_size = config.vocab_size
@@ -251,10 +262,6 @@ class LanguageModel(nn.Module):
         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__(
@@ -264,13 +271,29 @@ class LanguageModel(nn.Module):
     ) -> mx.array:
         h = self.embed_tokens(x)
 
-        if cache is None:
-            cache = [None] * self.num_layers
+        pipeline_rank = self.pipeline_rank
+        pipeline_size = self.pipeline_size
 
+        if cache is None:
+            cache = [None] * len(self.pipeline_layers)
         mask = create_attention_mask(h, cache[0])
 
-        for i in range(self.num_layers):
-            h = self.layers[self.start_idx + i](h, mask, cache[i])
+        # Receive from the previous process in the pipeline
+        if pipeline_rank < pipeline_size - 1:
+            h = mx.distributed.recv_like(h, (pipeline_rank + 1))
+
+        for l, c in zip(self.pipeline_layers, cache):
+            h = l(h, mask, cache=c)
+
+        # Send to the next process in the pipeline
+        if pipeline_rank != 0:
+            h = mx.distributed.send(h, (pipeline_rank - 1) % pipeline_size)
+            if cache[-1] is not None:
+                cache[-1].keys = mx.depends(cache[-1].keys, h)
+
+        # Broadcast h while keeping it in the graph
+        if pipeline_size > 1:
+            h = mx.distributed.all_gather(h)[: h.shape[0]]
 
         return self.norm(h)
 
@@ -313,9 +336,64 @@ class Model(nn.Module):
             if not k.startswith(f"model.layers.{mpt_layer}")
         }
 
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        for layer in self.model.layers:
+            # Shard the self attention
+            layer.self_attn.q_proj = shard_linear(
+                layer.self_attn.q_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.k_proj = shard_linear(
+                layer.self_attn.k_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.v_proj = shard_linear(
+                layer.self_attn.v_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.o_proj = shard_linear(
+                layer.self_attn.o_proj, "sharded-to-all", group=group
+            )
+            layer.self_attn.n_heads //= N
+            layer.self_attn.n_kv_heads //= N
+
+            # Shard the MLP
+            if isinstance(layer.mlp, MLP):
+                layer.mlp.gate_proj = shard_linear(
+                    layer.mlp.gate_proj, "all-to-sharded", group=group
+                )
+                layer.mlp.down_proj = shard_linear(
+                    layer.mlp.down_proj, "sharded-to-all", group=group
+                )
+                layer.mlp.up_proj = shard_linear(
+                    layer.mlp.up_proj, "all-to-sharded", group=group
+                )
+
+            # Shard the MoE. Shard in place since the MoE should be responsible
+            # for aggregating the results.
+            else:
+                layer.mlp.sharding_group = group
+                shard_inplace(
+                    layer.mlp.shared_experts.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_experts.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.shared_experts.up_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.up_proj, "all-to-sharded", group=group
+                )
+
     @property
     def layers(self):
-        return self.model.layers
+        return self.model.pipeline_layers
 
     @property
     def cast_predicate(self):

mlx_lm/models/glm4_moe_lite.py

@@ -0,0 +1,473 @@
+# Copyright © 2026 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 mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
+
+from .activations import swiglu
+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 = "glm4_moe_lite"
+    vocab_size: int = 154880
+    hidden_size: int = 2048
+    intermediate_size: int = 10240
+    moe_intermediate_size: int = 1536
+    num_hidden_layers: int = 47
+    num_attention_heads: int = 20
+    num_key_value_heads: int = 20
+    n_shared_experts: Optional[int] = 1
+    n_routed_experts: Optional[int] = 64
+    routed_scaling_factor: float = 1.8
+    kv_lora_rank: int = 512
+    q_lora_rank: int = 768
+    qk_rope_head_dim: int = 64
+    qk_nope_head_dim: int = 192
+    v_head_dim: int = 256
+    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 = 4
+    moe_layer_freq: int = 1
+    first_k_dense_replace: int = 1
+    max_position_embeddings: int = 202752
+    rms_norm_eps: float = 1e-5
+    rope_theta: float = 1_000_000.0
+    rope_scaling: Optional[Dict] = None
+    rope_traditional: bool = True
+    attention_bias: bool = False
+    attention_dropout: float = 0.0
+    partial_rotary_factor: float = 1.0
+    tie_word_embeddings: bool = False
+    num_nextn_predict_layers: int = 1
+
+
+class Glm4MoeLiteAttention(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=config.rms_norm_eps)
+            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=config.rms_norm_eps)
+        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=self.config.rope_traditional,
+            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 Glm4MoeLiteMLP(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(swiglu(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 Glm4MoeLiteMoE(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 = Glm4MoeLiteMLP(
+                config=config, intermediate_size=intermediate_size
+            )
+
+        self.sharding_group = None
+
+    def __call__(self, x):
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(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)
+
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+
+        return y
+
+
+class Glm4MoeLiteDecoderLayer(nn.Module):
+    def __init__(self, config: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.self_attn = Glm4MoeLiteAttention(config)
+        self.mlp = (
+            Glm4MoeLiteMoE(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 Glm4MoeLiteMLP(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 Glm4MoeLiteModel(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 = [
+            Glm4MoeLiteDecoderLayer(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
+        if pipeline_size > 1:
+            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 = Glm4MoeLiteModel(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):
+        # 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)
+
+        num_mpt_layers = getattr(self.args, "num_nextn_predict_layers", 0) or 0
+        if num_mpt_layers:
+
+            def _is_mpt_layer(key: str) -> bool:
+                for idx in range(num_mpt_layers):
+                    if key.startswith(
+                        f"model.layers.{self.args.num_hidden_layers + idx}"
+                    ):
+                        return True
+                return False
+
+            weights = {k: v for k, v in weights.items() if not _is_mpt_layer(k)}
+
+        return weights
+
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        for layer in self.model.layers:
+            # Shard the self attention
+            if layer.self_attn.q_lora_rank is None:
+                layer.self_attn.q_proj = shard_linear(
+                    layer.self_attn.q_proj, "all-to-sharded", group=group
+                )
+            else:
+                layer.self_attn.q_b_proj = shard_linear(
+                    layer.self_attn.q_b_proj, "all-to-sharded", group=group
+                )
+            layer.self_attn.kv_b_proj = shard_linear(
+                layer.self_attn.kv_b_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.o_proj = shard_linear(
+                layer.self_attn.o_proj, "sharded-to-all", group=group
+            )
+            layer.self_attn.num_heads //= N
+
+            # Shard the MLP
+            if isinstance(layer.mlp, Glm4MoeLiteMLP):
+                layer.mlp.gate_proj = shard_linear(
+                    layer.mlp.gate_proj, "all-to-sharded", group=group
+                )
+                layer.mlp.down_proj = shard_linear(
+                    layer.mlp.down_proj, "sharded-to-all", group=group
+                )
+                layer.mlp.up_proj = shard_linear(
+                    layer.mlp.up_proj, "all-to-sharded", group=group
+                )
+
+            # Shard the MoE. Shard in place since the MoE should be responsible
+            # for aggregating the results.
+            else:
+                layer.mlp.sharding_group = group
+                if getattr(layer.mlp, "shared_experts", None) is not None:
+                    shard_inplace(
+                        layer.mlp.shared_experts.gate_proj,
+                        "all-to-sharded",
+                        group=group,
+                    )
+                    shard_inplace(
+                        layer.mlp.shared_experts.down_proj,
+                        "sharded-to-all",
+                        group=group,
+                    )
+                    shard_inplace(
+                        layer.mlp.shared_experts.up_proj,
+                        "all-to-sharded",
+                        group=group,
+                    )
+                shard_inplace(
+                    layer.mlp.switch_mlp.gate_proj, "all-to-sharded", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.down_proj, "sharded-to-all", group=group
+                )
+                shard_inplace(
+                    layer.mlp.switch_mlp.up_proj, "all-to-sharded", group=group
+                )
+
+    @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,4 +1,4 @@
-# Copyright © 2023-2024 Apple Inc.
+# Copyright © 2023 - 2024 Apple Inc.
 
 from dataclasses import dataclass
 from typing import Any, Optional
@@ -138,7 +138,9 @@ class GPT2Model(nn.Module):
         if cache[0] is not None:
             offset = cache[0].offset
 
-        position_ids = mx.arange(offset, offset + L)
+        offset = mx.array(offset)
+        position_ids = mx.arange(L) + offset[..., None]
+
         hidden_states += self.wpe(position_ids)
 
         mask = create_attention_mask(hidden_states, cache[0])

mlx_lm/models/gpt_neox.py

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

mlx_lm/models/gpt_oss.py

@@ -7,6 +7,7 @@ from typing import Any, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
+from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .cache import KVCache, RotatingKVCache
@@ -142,8 +143,12 @@ class MLPBlock(nn.Module):
             bias=True,
         )
         self.router = nn.Linear(config.hidden_size, config.num_local_experts, bias=True)
+        self.sharding_group = None
 
     def __call__(self, x: mx.array) -> mx.array:
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(x)
+
         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)
@@ -152,7 +157,13 @@ class MLPBlock(nn.Module):
         x = self.experts(x, indices)
 
         x = x * mx.expand_dims(expert_weights, axis=-1)
-        return x.sum(axis=-2)
+
+        y = x.sum(axis=-2)
+
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+
+        return y
 
 
 class TransformerBlock(nn.Module):
@@ -268,6 +279,47 @@ class Model(nn.Module):
 
         return new_weights
 
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        R = group.rank()
+
+        for layer in self.model.layers:
+            layer.self_attn.q_proj = shard_linear(
+                layer.self_attn.q_proj, sharding="all-to-sharded", group=group
+            )
+            layer.self_attn.k_proj = shard_linear(
+                layer.self_attn.k_proj, sharding="all-to-sharded", group=group
+            )
+            layer.self_attn.v_proj = shard_linear(
+                layer.self_attn.v_proj, sharding="all-to-sharded", group=group
+            )
+            layer.self_attn.o_proj = shard_linear(
+                layer.self_attn.o_proj, sharding="sharded-to-all", group=group
+            )
+
+            layer.self_attn.num_attention_heads //= N
+            layer.self_attn.num_key_value_heads //= N
+            layer.self_attn.num_key_value_groups = (
+                layer.self_attn.num_attention_heads
+                // layer.self_attn.num_key_value_heads
+            )
+
+            layer.self_attn.sinks = layer.self_attn.sinks[
+                layer.self_attn.num_attention_heads
+                * R : layer.self_attn.num_attention_heads
+                * (R + 1)
+            ]
+
+            shard_inplace(layer.mlp.experts.gate_proj, "all-to-sharded", group=group)
+            shard_inplace(layer.mlp.experts.down_proj, "sharded-to-all", group=group)
+            layer.mlp.experts.down_proj.bias /= N
+            shard_inplace(
+                layer.mlp.experts.up_proj, sharding="all-to-sharded", group=group
+            )
+
+            layer.mlp.sharding_group = group
+
     @property
     def layers(self):
         return self.model.layers

mlx_lm/models/granite.py

@@ -6,6 +6,7 @@ from typing import Any, Dict, Optional, Union
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .rope_utils import initialize_rope
 
@@ -104,7 +105,7 @@ class MLP(nn.Module):
         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))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class TransformerBlock(nn.Module):

mlx_lm/models/granitemoehybrid.py

@@ -6,6 +6,7 @@ from typing import Any, List, Optional, Tuple
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import (
     BaseModelArgs,
     create_attention_mask,
@@ -20,6 +21,7 @@ from .switch_layers import SwitchGLU
 
 @dataclass
 class ModelArgs(BaseModelArgs):
+    # Required fields (no defaults)
     model_type: str
     vocab_size: int
     hidden_size: int
@@ -29,34 +31,42 @@ class ModelArgs(BaseModelArgs):
     num_attention_heads: int
     num_key_value_heads: int
     attention_bias: bool
-
-    # Scalar multipliers
     embedding_multiplier: float
     attention_multiplier: float
     logits_scaling: float
     residual_multiplier: float
-
-    # MoE parameters
-    num_local_experts: int
-    num_experts_per_tok: int
-    shared_intermediate_size: int
-
-    # Mamba parameters
-    mamba_n_heads: int
-    mamba_d_head: int
-    mamba_proj_bias: bool
-    mamba_d_state: int
-    mamba_d_conv: int
-    mamba_n_groups: int
-    mamba_conv_bias: bool
-
     layer_types: List[str]
     rms_norm_eps: float
     rope_theta: float
+
+    # Optional fields (with defaults)
+    # MoE parameters (optional for dense mode)
+    num_local_experts: Optional[int] = None
+    num_experts_per_tok: Optional[int] = None
+    shared_intermediate_size: Optional[int] = None
+
+    # Mamba parameters (optional for non-hybrid mode)
+    mamba_n_heads: Optional[int] = None
+    mamba_d_head: Optional[int] = None
+    mamba_proj_bias: Optional[bool] = None
+    mamba_d_state: Optional[int] = None
+    mamba_d_conv: Optional[int] = None
+    mamba_n_groups: Optional[int] = None
+    mamba_conv_bias: Optional[bool] = None
+
+    # Dense MLP parameters (for non-MoE mode)
+    mlp_bias: bool = False
+
+    # Other optional parameters
     position_embedding_type: str = "rope"
     tie_word_embeddings: bool = True
     time_step_limit: Tuple[float, float] = (0.001, 100.0)
 
+    # Mode flags - inferred from num_local_experts
+    @property
+    def use_moe(self) -> bool:
+        return bool(self.num_local_experts)
+
 
 class GraniteMoeHybridRMSNormGated(nn.Module):
     def __init__(self, hidden_size: int, eps: float = 1e-6):
@@ -66,7 +76,7 @@ class GraniteMoeHybridRMSNormGated(nn.Module):
 
     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)
+            hidden_states = swiglu(gate, hidden_states)
         return mx.fast.rms_norm(hidden_states, self.weight, self.eps)
 
 
@@ -110,21 +120,36 @@ class GraniteMoeHybridMamba2Mixer(nn.Module):
             self.intermediate_size, self.hidden_size, bias=args.mamba_proj_bias
         )
 
-    def _apply_conv(
-        self, conv_input: mx.array, cache: Optional[MambaCache] = None
+    def _conv(
+        self,
+        conv_input: mx.array,
+        cache: Optional[MambaCache],
+        mask: Optional[mx.array],
     ) -> 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 mask is not None:
+            conv_input = mx.where(mask[..., None], conv_input, 0)
 
         if cache is not None:
-            cache[0] = padded_input[:, -(self.conv_kernel_size - 1) :]
+            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)
+            n_keep = self.conv_kernel_size - 1
+            if cache.lengths is not None:
+                t = padded_input.shape[1]
+                ends = mx.clip(cache.lengths, 0, t - n_keep)
+                positions = (ends[:, None] + mx.arange(n_keep))[..., None]
+                cache[0] = mx.take_along_axis(padded_input, positions, axis=1)
+            else:
+                cache[0] = padded_input[:, -n_keep:, :]
+        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)
@@ -135,8 +160,8 @@ class GraniteMoeHybridMamba2Mixer(nn.Module):
         B: mx.array,
         C: mx.array,
         dt: mx.array,
-        state: Optional[mx.array] = None,
-        mask: Optional[mx.array] = None,
+        cache: Optional[MambaCache],
+        mask: Optional[mx.array],
     ) -> mx.array:
         batch_size, seq_len, _ = hidden_states.shape
 
@@ -145,26 +170,33 @@ class GraniteMoeHybridMamba2Mixer(nn.Module):
         )
         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)
+        if cache:
+            state = cache[1]
+            lengths = cache.lengths
+        else:
+            state, lengths = None, None
 
         y, state = ssm_update(
             hidden_states,
             self.A_log,
             B,
             C,
-            self.D,
+            self.D.astype(hidden_states.dtype),
             dt,
             self.dt_bias,
             state,
             self.time_step_limit,
             mask,
         )
+        if cache:
+            cache[1] = state
 
-        return y.reshape(batch_size, seq_len, self.intermediate_size), state
+        return y.reshape(batch_size, seq_len, self.intermediate_size)
 
     def __call__(
         self,
         hidden_states: mx.array,
-        mask: Optional[mx.array] = None,
+        mask: Optional[mx.array],
         cache: Optional[MambaCache] = None,
     ) -> mx.array:
 
@@ -175,11 +207,7 @@ class GraniteMoeHybridMamba2Mixer(nn.Module):
             [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)
-
+        conv_output = self._conv(conv_input, cache, mask)
         hidden_states_ssm, B, C = mx.split(
             conv_output,
             [
@@ -188,10 +216,9 @@ class GraniteMoeHybridMamba2Mixer(nn.Module):
             ],
             axis=-1,
         )
-        state = cache[1] if cache else None
-        y, state = self._ssm(hidden_states_ssm, B, C, dt, state, mask)
+        y = self._ssm(hidden_states_ssm, B, C, dt, cache, mask)
         if cache:
-            cache[1] = state
+            cache.advance(y.shape[1])
         y = self.norm(y, gate)
         return self.out_proj(y)
 
@@ -311,7 +338,22 @@ class GraniteMoeHybridSharedMLP(nn.Module):
 
     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)
+        return self.output_linear(swiglu(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(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class GraniteMoeHybridLayer(nn.Module):
@@ -319,6 +361,7 @@ class GraniteMoeHybridLayer(nn.Module):
         super().__init__()
         self.layer_type = layer_type
         self.residual_multiplier = args.residual_multiplier
+        self.use_moe = args.use_moe
 
         self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
 
@@ -329,8 +372,14 @@ class GraniteMoeHybridLayer(nn.Module):
         else:
             raise ValueError(f"Unknown layer type: {layer_type}")
 
-        self.shared_mlp = GraniteMoeHybridSharedMLP(args)
-        self.block_sparse_moe = GraniteMoeHybridMoE(args)
+        # MoE or dense MLP after attention/mamba
+        if self.use_moe:
+            self.shared_mlp = GraniteMoeHybridSharedMLP(args)
+            self.block_sparse_moe = GraniteMoeHybridMoE(args)
+        else:
+            # Dense MLP mode
+            self.mlp = GraniteMoeHybridMLP(args)
+
         self.post_attention_layernorm = nn.RMSNorm(
             args.hidden_size, eps=args.rms_norm_eps
         )
@@ -352,13 +401,16 @@ class GraniteMoeHybridLayer(nn.Module):
 
         hidden_states = residual + hidden_states * self.residual_multiplier
 
-        # Second block: MoE + shared_mlp
+        # Second block: MoE + shared_mlp OR dense MLP
         residual = hidden_states
         normed = self.post_attention_layernorm(hidden_states)
 
-        moe_out = self.block_sparse_moe(normed)
-        shared_out = self.shared_mlp(normed)
-        mlp_out = moe_out + shared_out
+        if self.use_moe:
+            moe_out = self.block_sparse_moe(normed)
+            shared_out = self.shared_mlp(normed)
+            mlp_out = moe_out + shared_out
+        else:
+            mlp_out = self.mlp(normed)
 
         hidden_states = residual + mlp_out * self.residual_multiplier
 
@@ -375,9 +427,16 @@ class GraniteMoeHybridModel(nn.Module):
         ]
         self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
         self.embedding_multiplier = args.embedding_multiplier
-        self.fa_idx = args.layer_types.index("attention")
-        self.ssm_idx = args.layer_types.index("mamba")
-        self.layer_types = args.layer_types
+
+        # Handle hybrid vs non-hybrid mode
+        self.fa_idx = (
+            args.layer_types.index("attention")
+            if "attention" in args.layer_types
+            else None
+        )
+        self.ssm_idx = (
+            args.layer_types.index("mamba") if "mamba" in args.layer_types else None
+        )
 
     def __call__(
         self,
@@ -389,11 +448,16 @@ class GraniteMoeHybridModel(nn.Module):
         if cache is None:
             cache = [None] * len(self.layers)
 
-        attn_mask = create_attention_mask(hidden_states, cache[self.fa_idx])
-        mamba_mask = create_ssm_mask(hidden_states, cache[self.ssm_idx])
+        # Create masks based on what layer types exist
+        attn_mask = None
+        mamba_mask = None
 
-        cache_counter = 0
-        for layer, c, layer_type in zip(self.layers, cache, self.layer_types):
+        if self.fa_idx is not None:
+            attn_mask = create_attention_mask(hidden_states, cache[self.fa_idx])
+        if self.ssm_idx is not None:
+            mamba_mask = create_ssm_mask(hidden_states, cache[self.ssm_idx])
+
+        for layer, c in zip(self.layers, cache):
             mask = attn_mask if layer.layer_type == "attention" else mamba_mask
             hidden_states = layer(hidden_states, mask=mask, cache=c)
 
@@ -443,8 +507,11 @@ class Model(nn.Module):
             if "conv1d.weight" in k and v.shape[-1] != 1:
                 weights[k] = v.moveaxis(2, 1)
 
-        # Handle MoE weight transformation to SwitchGLU format
-        if "model.layers.0.block_sparse_moe.input_linear.weight" in weights:
+        # Handle MoE weight transformation to SwitchGLU format (only for MoE models)
+        if (
+            self.args.use_moe
+            and "model.layers.0.block_sparse_moe.input_linear.weight" in weights
+        ):
             for l in range(self.args.num_hidden_layers):
                 prefix = f"model.layers.{l}.block_sparse_moe"
 
@@ -461,12 +528,31 @@ class Model(nn.Module):
                     f"{prefix}.output_linear.weight"
                 )
 
+        # Handle dense MLP weight transformation (for dense models)
+        elif (
+            not self.args.use_moe
+            and "model.layers.0.shared_mlp.input_linear.weight" in weights
+        ):
+            for l in range(self.args.num_hidden_layers):
+                prefix = f"model.layers.{l}.shared_mlp"
+
+                # Transform shared_mlp weights to standard mlp weights
+                input_weight = weights.pop(f"{prefix}.input_linear.weight")
+                # Split into gate and up projections (each half)
+                gate_proj, up_proj = mx.split(input_weight, 2, axis=0)
+                weights[f"model.layers.{l}.mlp.gate_proj.weight"] = gate_proj
+                weights[f"model.layers.{l}.mlp.up_proj.weight"] = up_proj
+
+                weights[f"model.layers.{l}.mlp.down_proj.weight"] = weights.pop(
+                    f"{prefix}.output_linear.weight"
+                )
+
         return weights
 
     @property
     def quant_predicate(self):
         def predicate(path, _):
-            if path.endswith("router.layer"):
+            if self.args.use_moe and path.endswith("router.layer"):
                 return {"group_size": 64, "bits": 8}
             return True
 

mlx_lm/models/helium.py

@@ -6,6 +6,7 @@ from typing import Any, Optional
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
@@ -92,7 +93,7 @@ class HeliumMLP(nn.Module):
         )
 
     def __call__(self, x: mx.array) -> mx.array:
-        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class HeliumDecoderLayer(nn.Module):

mlx_lm/models/hunyuan.py

@@ -6,6 +6,7 @@ from typing import Any, Dict, Optional, Tuple, Union
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .switch_layers import SwitchGLU
 
@@ -148,7 +149,7 @@ class MLP(nn.Module):
         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))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class Gate(nn.Module):

mlx_lm/models/hunyuan_v1_dense.py

@@ -6,6 +6,7 @@ from typing import Any, Dict, Optional, Union
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
@@ -144,7 +145,7 @@ class MLP(nn.Module):
         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))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class TransformerBlock(nn.Module):

mlx_lm/models/internlm2.py

@@ -6,6 +6,7 @@ from typing import Any, Dict, Optional, Union
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
@@ -156,7 +157,7 @@ class MLP(nn.Module):
         self.w3 = nn.Linear(dim, hidden_dim, bias=False)
 
     def __call__(self, x) -> mx.array:
-        return self.w2(nn.silu(self.w1(x)) * self.w3(x))
+        return self.w2(swiglu(self.w1(x), self.w3(x)))
 
 
 class TransformerBlock(nn.Module):

mlx_lm/models/internlm3.py

@@ -6,6 +6,7 @@ from typing import Any, Dict, Optional, Union
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
@@ -154,7 +155,7 @@ class MLP(nn.Module):
         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))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class TransformerBlock(nn.Module):

mlx_lm/models/iquestloopcoder.py

@@ -0,0 +1,286 @@
+# Copyright © 2026 Apple Inc.
+
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.nn.layers.distributed import shard_linear
+
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import KVCache, RotatingKVCache
+from .rope_utils import initialize_rope
+
+
+@partial(mx.compile, shapeless=True)
+def _compute_gate(query: mx.array, weight: mx.array, bias: mx.array) -> mx.array:
+    gate_logits = query @ weight[:, None, :].swapaxes(-1, -2)
+    gate_logits = gate_logits + bias[..., None, None]
+    return mx.sigmoid(gate_logits)
+
+
+@partial(mx.compile, shapeless=True)
+def _silu_mul(gate: mx.array, up: mx.array) -> mx.array:
+    return nn.silu(gate) * up
+
+
+@partial(mx.compile, shapeless=True)
+def _mix_attention(
+    gate: mx.array, attn_global: mx.array, attn_local: mx.array
+) -> mx.array:
+    return gate * attn_global + (1 - gate) * attn_local
+
+
+@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: int = 131072
+    attention_bias: bool = False
+    mlp_bias: bool = False
+    rope_theta: float = 500000.0
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = False
+    loop_num: int = 2
+    loop_window_size: int = 64
+
+
+class LoopGateProjection(nn.Module):
+    def __init__(self, num_heads: int, head_dim: int):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        self.weight = mx.zeros((num_heads, head_dim))
+        self.bias = mx.zeros((num_heads,))
+
+    def __call__(self, query: mx.array) -> mx.array:
+        return _compute_gate(query, self.weight, self.bias)
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.n_heads = n_heads = args.num_attention_heads
+        self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+        self.head_dim = head_dim = args.head_dim
+        self.scale = head_dim**-0.5
+
+        self.q_proj = nn.Linear(dim, n_heads * head_dim, bias=args.attention_bias)
+        self.k_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.attention_bias)
+        self.v_proj = nn.Linear(dim, n_kv_heads * head_dim, bias=args.attention_bias)
+        self.o_proj = nn.Linear(n_heads * head_dim, dim, bias=args.attention_bias)
+
+        self.rope = initialize_rope(
+            head_dim,
+            args.rope_theta,
+            traditional=False,
+            scaling_config=args.rope_scaling,
+            max_position_embeddings=args.max_position_embeddings,
+        )
+
+    def get_qkv(
+        self, x: mx.array, offset: int = 0
+    ) -> Tuple[mx.array, mx.array, mx.array]:
+        B, L, _ = x.shape
+        queries = self.q_proj(x).reshape(B, L, self.n_heads, -1).transpose(0, 2, 1, 3)
+        keys = self.k_proj(x).reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+        values = self.v_proj(x).reshape(B, L, self.n_kv_heads, -1).transpose(0, 2, 1, 3)
+
+        queries = self.rope(queries, offset=offset)
+        keys = self.rope(keys, offset=offset)
+
+        return queries, keys, values
+
+    def attention(
+        self,
+        queries: mx.array,
+        keys: mx.array,
+        values: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        return scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+
+
+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=args.mlp_bias)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=args.mlp_bias)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=args.mlp_bias)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(_silu_mul(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)
+        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
+        )
+
+
+class IQuestLoopCoderModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        assert args.loop_num == 2, f"Only loop_num=2 is supported, got {args.loop_num}"
+        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) for _ in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.gate_projections = [
+            LoopGateProjection(args.num_attention_heads, args.head_dim)
+            for _ in range(args.num_hidden_layers)
+        ]
+        self.loop_num = args.loop_num
+        self.loop_window_size = args.loop_window_size
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[List[Any]] = None,
+    ):
+        B, L = inputs.shape[:2]
+        h = self.embed_tokens(inputs)
+
+        if cache is None:
+            cache = [None] * (2 * len(self.layers))
+
+        mask = create_attention_mask(h, cache[0])
+        window_mask = create_attention_mask(
+            h, cache[len(self.layers)], window_size=self.loop_window_size
+        )
+
+        loop1_kv = []
+        for layer, c in zip(self.layers, cache):
+            h_norm = layer.input_layernorm(h)
+            offset = c.offset if c is not None else 0
+            q1, k1, v1 = layer.self_attn.get_qkv(h_norm, offset)
+
+            if c is not None:
+                k1, v1 = c.update_and_fetch(k1, v1)
+            loop1_kv.append((k1, v1))
+
+            out = layer.self_attn.attention(q1, k1, v1, mask, cache=c)
+            r = layer.self_attn.o_proj(out.transpose(0, 2, 1, 3).reshape(B, L, -1))
+            h = h + r
+            r = layer.mlp(layer.post_attention_layernorm(h))
+            h = h + r
+
+        for layer, gate_proj, c, (k1, v1) in zip(
+            self.layers, self.gate_projections, cache[len(self.layers) :], loop1_kv
+        ):
+            h_norm = layer.input_layernorm(h)
+            offset = c.offset if c is not None else 0
+            q2, k2, v2 = layer.self_attn.get_qkv(h_norm, offset)
+            gate = gate_proj(q2)
+            attn_global = layer.self_attn.attention(q2, k1, v1, mask, cache=c)
+
+            if c is not None:
+                k2, v2 = c.update_and_fetch(k2, v2)
+            attn_local = layer.self_attn.attention(
+                q2,
+                k2,
+                v2,
+                window_mask,
+                cache=c,
+            )
+
+            mixed = _mix_attention(gate, attn_global, attn_local)
+            r = layer.self_attn.o_proj(mixed.transpose(0, 2, 1, 3).reshape(B, L, -1))
+            h = h + r
+            r = layer.mlp(layer.post_attention_layernorm(h))
+            h = h + r
+
+        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 = IQuestLoopCoderModel(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 shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        rank = group.rank()
+
+        for i, layer in enumerate(self.model.layers):
+            layer.self_attn.q_proj = shard_linear(
+                layer.self_attn.q_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.k_proj = shard_linear(
+                layer.self_attn.k_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.v_proj = shard_linear(
+                layer.self_attn.v_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.o_proj = shard_linear(
+                layer.self_attn.o_proj, "sharded-to-all", group=group
+            )
+            layer.self_attn.n_heads //= N
+            layer.self_attn.n_kv_heads //= N
+
+            layer.mlp.gate_proj = shard_linear(
+                layer.mlp.gate_proj, "all-to-sharded", group=group
+            )
+            layer.mlp.down_proj = shard_linear(
+                layer.mlp.down_proj, "sharded-to-all", group=group
+            )
+            layer.mlp.up_proj = shard_linear(
+                layer.mlp.up_proj, "all-to-sharded", group=group
+            )
+
+            gate_proj = self.model.gate_projections[i]
+            heads_per_rank = gate_proj.num_heads // N
+            start = rank * heads_per_rank
+            end = start + heads_per_rank
+            gate_proj.weight = gate_proj.weight[start:end, :]
+            gate_proj.bias = gate_proj.bias[start:end]
+            gate_proj.num_heads = heads_per_rank
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def make_cache(self):
+        return [KVCache() for _ in self.layers] + [
+            RotatingKVCache(max_size=self.args.loop_window_size) for _ in self.layers
+        ]

mlx_lm/models/jamba.py

@@ -0,0 +1,385 @@
+# 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 .activations import swiglu
+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(swiglu(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(swiglu(z, y))
+        return z, (conv_state, ssm_state)
+
+    def __call__(self, x, cache):
+        if cache is None:
+            conv_state, ssm_state = None, None
+        else:
+            conv_state, ssm_state = cache[0], cache[1]
+
+        output, (conv_state, ssm_state) = self._process_sequence(
+            x, conv_state, ssm_state
+        )
+
+        if cache is not None:
+            cache[0] = conv_state
+            cache[1] = ssm_state
+
+        return output
+
+
+class JambaSparseMoeBlock(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_experts_per_tok = args.num_experts_per_tok
+
+        self.router = nn.Linear(args.hidden_size, args.num_experts, bias=False)
+        self.switch_mlp = SwitchGLU(
+            args.hidden_size, args.intermediate_size, args.num_experts
+        )
+
+    def __call__(self, x: mx.array) -> mx.array:
+        gates = self.router(x)
+        k = self.num_experts_per_tok
+        inds = mx.stop_gradient(mx.argpartition(-gates, kth=k - 1, axis=-1)[..., :k])
+        scores = mx.take_along_axis(gates, inds, axis=-1)
+        scores = mx.softmax(scores, axis=-1, precise=True)
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2)
+        return y
+
+
+class JambaDecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_type: str, layer_idx: int):
+        super().__init__()
+        self.is_attn = layer_type == "attention"
+        if self.is_attn:
+            self.self_attn = JambaAttention(args)
+        else:
+            self.mamba = JambaMambaMixer(args)
+        if (
+            args.num_experts > 1
+            and (layer_idx + args.expert_layer_offset) % args.expert_layer_period == 0
+        ):
+            ffn_layer_class = JambaSparseMoeBlock
+        else:
+            ffn_layer_class = JambaMLP
+        self.feed_forward = ffn_layer_class(args)
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.pre_ff_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        if self.is_attn:
+            h = self.self_attn(self.input_layernorm(x), mask, cache)
+        else:
+            h = self.mamba(self.input_layernorm(x), cache)
+        r = x + h
+        out = r + self.feed_forward(self.pre_ff_layernorm(r))
+        return out
+
+
+class JambaModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+
+        self.layers = [
+            JambaDecoderLayer(args, t, idx)
+            for idx, t in enumerate(args.layers_block_type)
+        ]
+        self.final_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.attn_idx = args.layers_block_type.index("attention")
+        self.ssm_idx = args.layers_block_type.index("mamba")
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        h = self.embed_tokens(inputs)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        attn_mask = create_attention_mask(h, cache[self.attn_idx])
+        ssm_mask = create_ssm_mask(h, cache[self.ssm_idx])
+
+        for layer, c in zip(self.layers, cache):
+            mask = attn_mask if layer.is_attn else ssm_mask
+            h = layer(h, mask=mask, cache=c)
+
+        return self.final_layernorm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.model_type = args.model_type
+        self.args = args
+        self.model = JambaModel(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        out = self.model(inputs, cache)
+        if self.args.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
+        return out
+
+    def make_cache(self):
+        caches = []
+        for layer in self.model.layers:
+            if layer.is_attn:
+                caches.append(KVCache())
+            else:
+                caches.append(MambaCache())
+        return caches
+
+    def sanitize(self, weights):
+        for k, v in list(weights.items()):
+            if "conv1d.weight" in k and v.shape[-1] != 1:
+                weights[k] = v.moveaxis(2, 1)
+
+        if self.args.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
+
+        for l in range(self.args.num_hidden_layers):
+            base = f"model.layers.{l}.feed_forward"
+            if not any(key.startswith(f"{base}.experts.") for key in weights.keys()):
+                continue
+
+            for proj in ["gate_proj", "down_proj", "up_proj"]:
+                for name in ["weight", "bias", "scales", "biases"]:
+                    expert_tensors = [
+                        weights.pop(f"{base}.experts.{e}.{proj}.{name}")
+                        for e in range(len(weights))
+                        if f"{base}.experts.{e}.{proj}.{name}" in weights
+                    ]
+                    if expert_tensors:
+                        weights[f"{base}.switch_mlp.{proj}.{name}"] = mx.stack(
+                            expert_tensors
+                        )
+
+        return weights
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if path.endswith("router"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate

mlx_lm/models/kimi_linear.py

@@ -0,0 +1,591 @@
+# 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 .activations import swiglu
+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(swiglu(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,
+        state: Optional[mx.array],
+        mask: Optional[mx.array],
+        lengths: Optional[mx.array],
+    ) -> Tuple[mx.array, mx.array]:
+        if mask is not None:
+            x = mx.where(mask[..., None], x, 0)
+
+        if state is None:
+            state = mx.zeros(
+                (x.shape[0], self.kernel_size - 1, x.shape[-1]), dtype=x.dtype
+            )
+        conv_input = mx.concatenate([state, x], axis=1)
+        out = nn.silu(self.conv(conv_input))
+        n_keep = self.kernel_size - 1
+        if lengths is not None:
+            ends = mx.clip(cache.lengths, 0, x.shape[1])
+            positions = (ends[:, None] + mx.arange(n_keep))[..., None]
+            new_state = mx.take_along_axis(conv_input, positions, axis=1)
+        else:
+            new_state = conv_input[:, -n_keep:, :]
+
+        return out, new_state
+
+
+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
+            lengths = cache.lengths
+        else:
+            conv_state = None
+            ssm_state = None
+            lengths = 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, mask, lengths)
+        k_conv, k_state = self.k_conv(self.k_proj(x), k_state, mask, lengths)
+        v_conv, v_state = self.v_conv(self.v_proj(x), v_state, mask, lengths)
+
+        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
+            cache.advance(T)
+
+        gate = self.g_b_proj(self.g_a_proj(x)).reshape(
+            B, T, self.num_heads, self.head_dim
+        )
+        out = (
+            self.o_norm(out.reshape(B, T, self.num_heads, self.head_dim))
+            * mx.sigmoid(gate)
+        ).reshape(B, T, -1)
+        return self.o_proj(out)
+
+
+class KimiDecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        kda_layers = args.linear_attn_config["kda_layers"]
+        self.is_linear = (layer_idx + 1) in kda_layers
+
+        if self.is_linear:
+            self.self_attn = KimiDeltaAttention(args, layer_idx)
+        else:
+            self.self_attn = KimiMLAAttention(args)
+
+        if (
+            args.num_experts > 0
+            and layer_idx >= args.first_k_dense_replace
+            and layer_idx % args.moe_layer_freq == 0
+        ):
+            self.mlp = KimiSparseMoE(args)
+        else:
+            self.mlp = KimiMLP(args)
+
+        self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        attn_cache = None if cache is None else cache
+        y = self.self_attn(self.input_layernorm(x), mask, attn_cache)
+        h = x + y
+        z = self.mlp(self.post_attention_layernorm(h))
+        return h + z
+
+
+class KimiLinearModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [KimiDecoderLayer(args, i) for i in range(args.num_hidden_layers)]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        kda_layers = args.linear_attn_config["kda_layers"]
+        self.ssm_idx = kda_layers[0] - 1
+        for i in range(len(self.layers)):
+            if (i + 1) not in kda_layers:
+                self.attn_idx = i
+                break
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[List[Any]] = None,
+    ) -> mx.array:
+        h = self.embed_tokens(inputs)
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        ssm_mask = create_ssm_mask(h, cache[self.ssm_idx])
+        attn_mask = create_attention_mask(h, cache[self.attn_idx])
+
+        for layer, layer_cache in zip(self.layers, cache):
+            mask = ssm_mask if layer.is_linear else attn_mask
+            h = layer(h, mask=mask, cache=layer_cache)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = KimiLinearModel(args)
+        if args.tie_word_embeddings:
+            self.lm_head = None
+        else:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[List[Any]] = None,
+    ) -> mx.array:
+        out = self.model(inputs, cache)
+        if self.lm_head is None:
+            return self.model.embed_tokens.as_linear(out)
+        return self.lm_head(out)
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def make_cache(self):
+        caches: List[Any] = []
+        for layer in self.layers:
+            if layer.is_linear:
+                caches.append(MambaCache())
+            else:
+                caches.append(KVCache())
+        return caches
+
+    def sanitize(self, weights: Dict[str, mx.array]) -> Dict[str, mx.array]:
+        weights = {k: v for k, v in weights.items() if not k.startswith("model.mtp")}
+
+        if self.args.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
+
+        for layer_idx, layer in enumerate(self.layers):
+            prefix = f"model.layers.{layer_idx}"
+
+            if isinstance(layer.mlp, KimiSparseMoE):
+                src_prefix = f"{prefix}.block_sparse_moe"
+                dst_prefix = f"{prefix}.mlp"
+                for src, dst in [
+                    ("w1", "gate_proj"),
+                    ("w2", "down_proj"),
+                    ("w3", "up_proj"),
+                ]:
+                    key = f"{src_prefix}.experts.0.{src}.weight"
+                    if key in weights:
+                        stacked = [
+                            weights.pop(f"{src_prefix}.experts.{i}.{src}.weight")
+                            for i in range(self.args.num_experts)
+                        ]
+                        weights[f"{dst_prefix}.switch_mlp.{dst}.weight"] = mx.stack(
+                            stacked
+                        )
+
+                for name in ("gate_proj", "up_proj", "down_proj"):
+                    src_key = f"{src_prefix}.shared_experts.{name}.weight"
+                    if src_key in weights:
+                        weights[f"{dst_prefix}.shared_experts.{name}.weight"] = (
+                            weights.pop(src_key)
+                        )
+
+                gate_key = f"{src_prefix}.gate.weight"
+                if gate_key in weights:
+                    weights[f"{dst_prefix}.gate.weight"] = weights.pop(gate_key)
+
+                bias_key = f"{src_prefix}.gate.e_score_correction_bias"
+                if bias_key in weights:
+                    weights[f"{dst_prefix}.e_score_correction_bias"] = weights.pop(
+                        bias_key
+                    )
+
+            attn = getattr(layer, "self_attn", None)
+            if isinstance(attn, KimiDeltaAttention):
+                attn_prefix = f"{prefix}.self_attn"
+                for src_name, dst_name in (
+                    ("q_conv1d", "q_conv"),
+                    ("k_conv1d", "k_conv"),
+                    ("v_conv1d", "v_conv"),
+                ):
+                    src_key = f"{attn_prefix}.{src_name}.weight"
+                    if src_key in weights:
+                        w = weights.pop(src_key)
+                        if w.ndim == 3:
+                            w = w.moveaxis(2, 1)
+                        weights[f"{attn_prefix}.{dst_name}.conv.weight"] = w
+                dt_key = f"{attn_prefix}.dt_bias"
+                if dt_key in weights:
+                    if weights[dt_key].ndim > 1:
+                        weights[dt_key] = mx.reshape(weights[dt_key], (-1,))
+
+        return weights
+
+    @property
+    def cast_predicate(self):
+        def predicate(path: str):
+            if "e_score_correction_bias" in path:
+                return False
+            if path.endswith("A_log") or path.endswith("dt_bias"):
+                return False
+            return True
+
+        return predicate
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if path.endswith("mlp.gate"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate

mlx_lm/models/lfm2-vl.py

@@ -18,11 +18,6 @@ class ModelArgs(BaseModelArgs):
 
     def __post_init__(self):
         self.text_config["tie_word_embeddings"] = False
-        self.text_config["full_attn_idxs"] = [
-            i
-            for i, layer_type in enumerate(self.text_config["layer_types"])
-            if layer_type == "full_attention"
-        ]
 
 
 class Model(nn.Module):

mlx_lm/models/lfm2.py

@@ -5,6 +5,7 @@ from typing import Any, List, Optional
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import (
     BaseModelArgs,
     create_attention_mask,
@@ -31,8 +32,19 @@ class ModelArgs(BaseModelArgs):
     block_multiple_of: int
     block_ffn_dim_multiplier: float
     block_auto_adjust_ff_dim: bool
-    full_attn_idxs: List[int]
     rope_theta: float
+    full_attn_idxs: Optional[List[int]] = None
+    layer_types: Optional[List[str]] = None
+
+    def __post_init__(self):
+        if self.num_key_value_heads is None:
+            self.num_key_value_heads = self.num_attention_heads
+        if self.full_attn_idxs is None:
+            self.full_attn_idxs = [
+                i
+                for i, layer_type in enumerate(self.layer_types)
+                if layer_type == "full_attention"
+            ]
 
 
 class Attention(nn.Module):
@@ -122,23 +134,33 @@ class ShortConv(nn.Module):
         mask: Optional[mx.array] = None,
         cache: Optional[Any] = None,
     ):
-        seqlen = x.shape[1]
         BCx = self.in_proj(x)
         B, C, x = mx.split(BCx, 3, axis=-1)
         Bx = B * x
         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) :]
+            if cache[0] is None:
+                state = mx.zeros(
+                    (Bx.shape[0], self.L_cache - 1, self.args.hidden_size),
+                    dtype=Bx.dtype,
+                )
+            else:
+                state = cache[0]
+            Bx = mx.concatenate([state, Bx], axis=1)
+            n_keep = self.L_cache - 1
+            t = x.shape[1]
+            if cache.lengths is not None:
+                ends = mx.clip(cache.lengths, 0, t)
+                positions = (ends[:, None] + mx.arange(n_keep))[..., None]
+                cache[0] = mx.take_along_axis(Bx, positions, axis=1)
+            else:
+                cache[0] = Bx[:, -n_keep:, :]
+            cache.advance(t)
+        else:
+            Bx = mx.pad(Bx, [(0, 0), (self.L_cache - 1, 0), (0, 0)])
+
         conv_out = self.conv(Bx)
 
         y = C * conv_out
@@ -166,7 +188,7 @@ class MLP(nn.Module):
         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))
+        return self.w2(swiglu(self.w1(x), self.w3(x)))
 
 
 class Lfm2DecoderLayer(nn.Module):

mlx_lm/models/lfm2_moe.py

@@ -0,0 +1,384 @@
+# 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 .activations import swiglu
+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)
+
+        if cache is not None:
+            if cache[0] is None:
+                state = mx.zeros(
+                    (Bx.shape[0], self.L_cache - 1, self.args.hidden_size),
+                    dtype=Bx.dtype,
+                )
+            else:
+                state = cache[0]
+            Bx = mx.concatenate([state, Bx], axis=1)
+            n_keep = self.L_cache - 1
+            t = x.shape[1]
+            if cache.lengths is not None:
+                ends = mx.clip(cache.lengths, 0, t)
+                positions = (ends[:, None] + mx.arange(n_keep))[..., None]
+                cache[0] = mx.take_along_axis(Bx, positions, axis=1)
+            else:
+                cache[0] = Bx[:, -n_keep:, :]
+            cache.advance(t)
+        else:
+            Bx = mx.pad(Bx, [(0, 0), (self.L_cache - 1, 0), (0, 0)])
+
+        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(swiglu(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

@@ -6,6 +6,7 @@ from typing import Any, Optional
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
@@ -87,7 +88,7 @@ class Lille130mMLP(nn.Module):
 
     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))
+        return self.down_proj(swiglu(self.gate_proj(h), self.up_proj(h)))
 
 
 class Lille130Block(nn.Module):

mlx_lm/models/llama.py

@@ -1,12 +1,15 @@
 # Copyright © 2023-2024 Apple Inc.
 
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 import mlx.core as mx
 import mlx.nn as nn
+from mlx.nn.layers.distributed import shard_linear
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import KVCache, RotatingKVCache
 from .rope_utils import initialize_rope
 
 
@@ -28,11 +31,16 @@ class ModelArgs(BaseModelArgs):
     rope_traditional: bool = False
     rope_scaling: Optional[Dict[str, Union[float, str]]] = None
     tie_word_embeddings: bool = True
+    layer_types: Optional[List[str]] = None
+    sliding_window: Optional[int] = None
 
     def __post_init__(self):
         if self.num_key_value_heads is None:
             self.num_key_value_heads = self.num_attention_heads
 
+        if self.layer_types is None:
+            self.layer_types = ["full_attention"] * self.num_hidden_layers
+
 
 class Attention(nn.Module):
     def __init__(self, args: ModelArgs):
@@ -110,14 +118,15 @@ class MLP(nn.Module):
         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))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class TransformerBlock(nn.Module):
-    def __init__(self, args: ModelArgs):
+    def __init__(self, args: ModelArgs, use_sliding: bool = False):
         super().__init__()
         self.num_attention_heads = args.num_attention_heads
         self.hidden_size = args.hidden_size
+        self.use_sliding = use_sliding
         self.self_attn = Attention(args)
         self.mlp = MLP(args)
         self.input_layernorm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
@@ -145,12 +154,21 @@ class LlamaModel(nn.Module):
         self.args = args
         self.vocab_size = args.vocab_size
         self.num_hidden_layers = args.num_hidden_layers
+        self.layer_types = args.layer_types
+        self.sliding_window = args.sliding_window
         assert self.vocab_size > 0
         self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
         self.layers = [
-            TransformerBlock(args=args) for _ in range(args.num_hidden_layers)
+            TransformerBlock(args=args, use_sliding=layer_type == "sliding_attention")
+            for layer_type in self.layer_types
         ]
         self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.fa_idx = self.layer_types.index("full_attention")
+        self.swa_idx = None
+        for e, l in enumerate(self.layers):
+            if l.use_sliding:
+                self.swa_idx = e
+                break
 
     def __call__(
         self,
@@ -166,10 +184,15 @@ class LlamaModel(nn.Module):
         if cache is None:
             cache = [None] * len(self.layers)
 
-        mask = create_attention_mask(h, cache[0])
+        fa_mask = create_attention_mask(h, cache[self.fa_idx])
+        if self.swa_idx is not None:
+            swa_mask = create_attention_mask(
+                h, cache[self.swa_idx], window_size=self.sliding_window
+            )
 
-        for layer, c in zip(self.layers, cache):
-            h = layer(h, mask, cache=c)
+        for layer, cache in zip(self.layers, cache):
+            mask = swa_mask if layer.use_sliding else fa_mask
+            h = layer(h, mask, cache=cache)
 
         return self.norm(h)
 
@@ -205,6 +228,47 @@ class Model(nn.Module):
             weights.pop("lm_head.weight", None)
         return weights
 
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        for layer in self.model.layers:
+            # Shard the self attention
+            layer.self_attn.q_proj = shard_linear(
+                layer.self_attn.q_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.k_proj = shard_linear(
+                layer.self_attn.k_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.v_proj = shard_linear(
+                layer.self_attn.v_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.o_proj = shard_linear(
+                layer.self_attn.o_proj, "sharded-to-all", group=group
+            )
+            layer.self_attn.n_heads //= N
+            layer.self_attn.n_kv_heads //= N
+
+            # Shard the MLP
+            layer.mlp.gate_proj = shard_linear(
+                layer.mlp.gate_proj, "all-to-sharded", group=group
+            )
+            layer.mlp.down_proj = shard_linear(
+                layer.mlp.down_proj, "sharded-to-all", group=group
+            )
+            layer.mlp.up_proj = shard_linear(
+                layer.mlp.up_proj, "all-to-sharded", group=group
+            )
+
     @property
     def layers(self):
         return self.model.layers
+
+    def make_cache(self):
+        return [
+            (
+                RotatingKVCache(max_size=self.model.sliding_window)
+                if layer.use_sliding
+                else KVCache()
+            )
+            for layer in self.layers
+        ]

mlx_lm/models/llama4.py

@@ -6,6 +6,7 @@ from typing import Any, Optional, Union
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .cache import ChunkedKVCache, KVCache
 from .rope_utils import initialize_rope
@@ -145,7 +146,7 @@ class MLP(nn.Module):
         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))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class MoE(nn.Module):

mlx_lm/models/llama4_text.py

@@ -6,6 +6,7 @@ from typing import Any, Optional
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
@@ -95,7 +96,7 @@ class MLP(nn.Module):
         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))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class TransformerBlock(nn.Module):

mlx_lm/models/longcat_flash.py

@@ -4,9 +4,12 @@ from typing import Any, Dict, Optional, Tuple
 
 import mlx.core as mx
 import mlx.nn as nn
+from mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
 
+from .activations import swiglu
 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
 
 
@@ -38,6 +41,7 @@ class ModelArgs(BaseModelArgs):
     attention_bias: bool
     norm_topk_prob: bool = False
     router_bias: bool = False
+    rope_scaling: Optional[Dict] = None
 
 
 class LongcatFlashMLA(nn.Module):
@@ -93,8 +97,20 @@ class LongcatFlashMLA(nn.Module):
         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
+        if args.rope_scaling is not None:
+            mscale_all_dim = args.rope_scaling.get("mscale_all_dim", 0)
+            if mscale_all_dim:
+                scaling_factor = args.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=args.rope_theta,
+            traditional=True,
+            scaling_config=args.rope_scaling,
+            max_position_embeddings=args.max_position_embeddings,
         )
 
     def __call__(
@@ -168,7 +184,7 @@ class LongcatFlashMLP(nn.Module):
         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))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class LongcatFlashTopkRouter(nn.Module):
@@ -223,8 +239,11 @@ class LongcatFlashMoE(nn.Module):
         )
 
         self.router = LongcatFlashTopkRouter(args)
+        self.sharding_group = None
 
     def __call__(self, hidden_states):
+        if self.sharding_group is not None:
+            hidden_states = sum_gradients(self.sharding_group)(hidden_states)
 
         topk_indices, topk_weights = self.router(hidden_states)
 
@@ -235,15 +254,21 @@ class LongcatFlashMoE(nn.Module):
 
         regular_outputs = self.switch_mlp(hidden_states, topk_indices)
 
-        weighted_outputs = regular_outputs * topk_weights[..., None]
-
-        # Add identity expert contribution if needed
-        assert self.zero_expert_type == "identity"
-        identity_weights = mx.where(mask, topk_weights, 0.0)
-        identity_outputs = hidden_states[..., None, :] * identity_weights[..., None]
-        weighted_outputs = weighted_outputs + identity_outputs
-
+        weighted_outputs = regular_outputs * regular_weights[..., None]
         final_output = mx.sum(weighted_outputs, axis=-2)
+
+        if self.sharding_group is not None:
+            final_output = mx.distributed.all_sum(
+                final_output, group=self.sharding_group
+            )
+
+        # Add identity expert contribution after all_sum to avoid summing it N times
+        assert self.zero_expert_type == "identity"
+        identity_weights_sum = mx.sum(
+            mx.where(mask, topk_weights, 0.0), axis=-1, keepdims=True
+        )
+        final_output = final_output + hidden_states * identity_weights_sum
+
         return final_output
 
 
@@ -379,3 +404,37 @@ class Model(nn.Module):
 
     def make_cache(self):
         return [CacheList(KVCache(), KVCache()) for _ in self.model.layers]
+
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+
+        for layer in self.model.layers:
+            for attn in layer.self_attn:
+                if attn.q_lora_rank is None:
+                    attn.q_proj = shard_linear(
+                        attn.q_proj, "all-to-sharded", group=group
+                    )
+                else:
+                    attn.q_b_proj = shard_linear(
+                        attn.q_b_proj, "all-to-sharded", group=group
+                    )
+                attn.kv_b_proj = shard_linear(
+                    attn.kv_b_proj, "all-to-sharded", group=group
+                )
+                attn.o_proj = shard_linear(attn.o_proj, "sharded-to-all", group=group)
+                attn.num_attention_heads //= N
+
+            for mlp in layer.mlps:
+                mlp.gate_proj = shard_linear(
+                    mlp.gate_proj, "all-to-sharded", group=group
+                )
+                mlp.up_proj = shard_linear(mlp.up_proj, "all-to-sharded", group=group)
+                mlp.down_proj = shard_linear(
+                    mlp.down_proj, "sharded-to-all", group=group
+                )
+
+            layer.mlp.sharding_group = group
+            shard_inplace(layer.mlp.switch_mlp.gate_proj, "all-to-sharded", group=group)
+            shard_inplace(layer.mlp.switch_mlp.up_proj, "all-to-sharded", group=group)
+            shard_inplace(layer.mlp.switch_mlp.down_proj, "sharded-to-all", group=group)

mlx_lm/models/mamba.py

@@ -6,6 +6,7 @@ from dataclasses import dataclass
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs
 from .cache import MambaCache
 
@@ -139,7 +140,7 @@ class MambaBlock(nn.Module):
             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)
+        z = self.out_proj(swiglu(z, y))
         return z, (new_conv_cache, current_state)
 
     def __call__(self, x, cache):

mlx_lm/models/mamba2.py

@@ -7,6 +7,7 @@ from typing import Optional, Tuple, Union
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_ssm_mask
 from .cache import MambaCache
 from .ssm import ssm_update
@@ -48,7 +49,7 @@ class MambaRMSNormGated(nn.Module):
 
     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)
+            hidden_states = swiglu(gate, hidden_states)
         return mx.fast.rms_norm(hidden_states, self.weight, self.eps)
 
 
@@ -93,9 +94,15 @@ class Mamba2Block(nn.Module):
             self.intermediate_size, self.hidden_size, bias=args.use_bias
         )
 
-    def _apply_conv(
-        self, conv_input: mx.array, cache: Optional[MambaCache] = None
+    def _conv(
+        self,
+        conv_input: mx.array,
+        cache: Optional[MambaCache],
+        mask: Optional[mx.array],
     ) -> mx.array:
+        if mask is not None:
+            conv_input = mx.where(mask[..., None], conv_input, 0)
+
         if cache is not None:
             if cache[0] is None:
                 conv_state = mx.zeros(
@@ -105,7 +112,14 @@ class Mamba2Block(nn.Module):
             else:
                 conv_state = cache[0]
             padded_input = mx.concatenate([conv_state, conv_input], axis=1)
-            cache[0] = padded_input[:, -(self.conv_kernel_size - 1) :, :]
+            n_keep = self.conv_kernel_size - 1
+            if cache.lengths is not None:
+                t = padded_input.shape[1]
+                ends = mx.clip(cache.lengths, 0, t - n_keep)
+                positions = (ends[:, None] + mx.arange(n_keep))[..., None]
+                cache[0] = mx.take_along_axis(padded_input, positions, axis=1)
+            else:
+                cache[0] = padded_input[:, -n_keep:, :]
         else:
             padded_input = mx.pad(
                 conv_input, [(0, 0), (self.conv_kernel_size - 1, 0), (0, 0)]
@@ -120,8 +134,8 @@ class Mamba2Block(nn.Module):
         B: mx.array,
         C: mx.array,
         dt: mx.array,
-        state: Optional[mx.array] = None,
-        mask: Optional[mx.array] = None,
+        cache: Optional[MambaCache],
+        mask: Optional[mx.array],
     ) -> mx.array:
         batch_size, seq_len, _ = hidden_states.shape
         hidden_states = hidden_states.reshape(
@@ -129,6 +143,11 @@ class Mamba2Block(nn.Module):
         )
         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)
+        if cache:
+            state = cache[1]
+            lengths = cache.lengths
+        else:
+            state, lengths = None, None
         y, state = ssm_update(
             hidden_states,
             self.A_log,
@@ -140,8 +159,11 @@ class Mamba2Block(nn.Module):
             state,
             self.time_step_limit,
             mask,
+            lengths,
         )
-        return y.reshape(batch_size, seq_len, self.intermediate_size), state
+        if cache:
+            cache[1] = state
+        return y.reshape(batch_size, seq_len, self.intermediate_size)
 
     def __call__(
         self,
@@ -155,9 +177,7 @@ class Mamba2Block(nn.Module):
             [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)
+        conv_output = self._conv(conv_input, cache, mask)
         hidden_states, B, C = mx.split(
             conv_output,
             [
@@ -166,10 +186,9 @@ class Mamba2Block(nn.Module):
             ],
             axis=-1,
         )
-        state = cache[1] if cache else None
-        y, state = self._ssm(hidden_states, B, C, dt, state, mask=mask)
+        y = self._ssm(hidden_states, B, C, dt, cache, mask=mask)
         if cache:
-            cache[1] = state
+            cache.advance(y.shape[1])
         y = self.norm(y, gate)
         return self.out_proj(y)
 

mlx_lm/models/mimo.py

@@ -6,6 +6,7 @@ from typing import Any, Dict, Optional, Union
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .rope_utils import initialize_rope
 
@@ -90,7 +91,7 @@ class MLP(nn.Module):
         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))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class TransformerBlock(nn.Module):

mlx_lm/models/mimo_v2_flash.py

@@ -0,0 +1,384 @@
+# Copyright © 2024 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 .activations import swiglu
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import KVCache, RotatingKVCache
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    num_experts_per_tok: int
+    hybrid_layer_pattern: List[int]
+    moe_layer_freq: List[int]
+    add_swa_attention_sink_bias: bool
+    add_full_attention_sink_bias: bool
+    sliding_window_size: int
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    moe_intermediate_size: int
+    num_hidden_layers: int
+    num_attention_heads: int
+    num_key_value_heads: int
+    n_shared_experts: Optional[int]
+    n_routed_experts: Optional[int]
+    routed_scaling_factor: Optional[float]
+    topk_method: str
+    scoring_func: str
+    norm_topk_prob: bool
+    n_group: int
+    topk_group: int
+    max_position_embeddings: int
+    layernorm_epsilon: float
+    rope_theta: float
+    swa_rope_theta: float
+    swa_num_attention_heads: int
+    swa_num_key_value_heads: int
+    head_dim: int
+    v_head_dim: int
+    swa_head_dim: int
+    swa_v_head_dim: int
+    partial_rotary_factor: int
+
+
+class Attention(nn.Module):
+    def __init__(self, args: ModelArgs, is_sliding_window: bool):
+        super().__init__()
+
+        dim = args.hidden_size
+        self.is_sliding_window = is_sliding_window
+        if self.is_sliding_window:
+            self.n_heads = n_heads = args.swa_num_attention_heads
+            self.n_kv_heads = n_kv_heads = args.swa_num_key_value_heads
+            self.has_sinks = args.add_swa_attention_sink_bias
+            head_dim = args.swa_head_dim
+            v_head_dim = args.swa_v_head_dim
+            rope_theta = args.swa_rope_theta
+        else:
+            self.n_heads = n_heads = args.num_attention_heads
+            self.n_kv_heads = n_kv_heads = args.num_key_value_heads
+            self.has_sinks = args.add_full_attention_sink_bias
+            head_dim = args.head_dim
+            v_head_dim = args.v_head_dim
+            rope_theta = args.rope_theta
+
+        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 * v_head_dim, bias=False)
+        self.o_proj = nn.Linear(n_heads * v_head_dim, dim, bias=False)
+        if self.has_sinks:
+            self.attention_sink_bias = mx.ones((self.n_heads,))
+        else:
+            self.attention_sink_bias = None
+
+        self.rope = nn.RoPE(
+            int(args.partial_rotary_factor * head_dim),
+            traditional=False,
+            base=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)
+
+        output = scaled_dot_product_attention(
+            queries,
+            keys,
+            values,
+            cache=cache,
+            scale=self.scale,
+            mask=mask,
+            sinks=self.attention_sink_bias,
+        )
+        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(swiglu(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 + 1e-20)
+    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
+            if config.routed_scaling_factor is not None
+            else 1.0
+        )
+        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, is_moe, is_sliding_window):
+        super().__init__()
+        self.self_attn = Attention(config, is_sliding_window)
+        self.mlp = MoE(config) if is_moe else MLP(config)
+        self.is_sliding_window = is_sliding_window
+        self.input_layernorm = nn.RMSNorm(
+            config.hidden_size, eps=config.layernorm_epsilon
+        )
+        self.post_attention_layernorm = nn.RMSNorm(
+            config.hidden_size, eps=config.layernorm_epsilon
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.self_attn(self.input_layernorm(x), mask, cache)
+        h = x + r
+        r = self.mlp(self.post_attention_layernorm(h))
+        return h + r
+
+
+class LanguageModel(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.layers = [
+            DecoderLayer(
+                config,
+                is_moe=config.moe_layer_freq[idx] == 1,
+                is_sliding_window=config.hybrid_layer_pattern[idx] == 1,
+            )
+            for idx in range(config.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(config.hidden_size, eps=config.layernorm_epsilon)
+        self.swa_idx = config.hybrid_layer_pattern.index(1)
+        self.ga_idx = config.hybrid_layer_pattern.index(0)
+        self.sliding_window_size = config.sliding_window_size
+
+    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)
+
+        full_mask = create_attention_mask(x, cache[self.ga_idx])
+        swa_mask = create_attention_mask(
+            x, cache[self.swa_idx], window_size=self.sliding_window_size
+        )
+
+        for l, c in zip(self.layers, cache):
+            mask = swa_mask if l.is_sliding_window else full_mask
+            h = l(h, mask, cache=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 = 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):
+        def dequant(weight, scale_inv):
+            dtype = mx.bfloat16
+            weight = mx.from_fp8(weight, dtype=mx.bfloat16)
+            bs = 128  # block size
+            m, n = weight.shape
+            pad_bottom = bs * scale_inv.shape[0] - m
+            pad_side = bs * scale_inv.shape[1] - n
+            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 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
+        return {k: v for k, v in weights.items() if not k.startswith("model.mtp")}
+
+    @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
+
+    def make_cache(self):
+        caches = []
+        for l in self.layers:
+            if l.is_sliding_window:
+                caches.append(RotatingKVCache(max_size=self.args.sliding_window_size))
+            else:
+                caches.append(KVCache())
+        return caches

mlx_lm/models/minicpm.py

@@ -6,6 +6,7 @@ from typing import Any, Dict, Optional, Union
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .rope_utils import initialize_rope
 
@@ -38,7 +39,7 @@ class MLP(nn.Module):
         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))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class Attention(nn.Module):

mlx_lm/models/minicpm3.py

@@ -6,6 +6,7 @@ from typing import Any, Dict, Optional, Union
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .rope_utils import SuScaledRoPE
 
@@ -156,7 +157,7 @@ class MLP(nn.Module):
         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))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class DecoderLayer(nn.Module):

mlx_lm/models/minimax.py

@@ -0,0 +1,344 @@
+# 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 mlx.nn.layers.distributed import shard_inplace, shard_linear, sum_gradients
+
+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,))
+        self.sharding_group = None
+
+    def __call__(self, x: mx.array) -> mx.array:
+        if self.sharding_group is not None:
+            x = sum_gradients(self.sharding_group)(x)
+
+        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)
+
+        if self.sharding_group is not None:
+            y = mx.distributed.all_sum(y, group=self.sharding_group)
+
+        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 = mx.bfloat16
+            weight = mx.from_fp8(weight, dtype=mx.bfloat16)
+            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
+
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        rank = group.rank()
+        for layer in self.model.layers:
+            # Shard the self attention
+            layer.self_attn.q_proj = shard_linear(
+                layer.self_attn.q_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.k_proj = shard_linear(
+                layer.self_attn.k_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.v_proj = shard_linear(
+                layer.self_attn.v_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.o_proj = shard_linear(
+                layer.self_attn.o_proj, "sharded-to-all", group=group
+            )
+            if layer.self_attn.use_qk_norm:
+                layer.self_attn.q_norm.weight = layer.self_attn.q_norm.weight.split(
+                    N, axis=-1
+                )[rank]
+                layer.self_attn.k_norm.weight = layer.self_attn.k_norm.weight.split(
+                    N, axis=-1
+                )[rank]
+
+            layer.self_attn.num_attention_heads //= N
+            layer.self_attn.num_key_value_heads //= N
+
+            # Shard the MLP
+            shard_inplace(
+                layer.block_sparse_moe.switch_mlp.gate_proj,
+                "all-to-sharded",
+                group=group,
+            )
+            shard_inplace(
+                layer.block_sparse_moe.switch_mlp.down_proj,
+                "sharded-to-all",
+                group=group,
+            )
+            shard_inplace(
+                layer.block_sparse_moe.switch_mlp.up_proj,
+                "all-to-sharded",
+                group=group,
+            )
+            layer.block_sparse_moe.sharding_group = group
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    @property
+    def cast_predicate(self):
+        def predicate(k):
+            return "e_score_correction_bias" not in k
+
+        return predicate
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if path.endswith("block_sparse_moe.gate"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate

mlx_lm/models/ministral3.py

@@ -0,0 +1,333 @@
+# 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 mlx.nn.layers.distributed import shard_linear
+
+from .activations import swiglu
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import KVCache, RotatingKVCache
+from .pipeline import PipelineMixin
+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(size, offset, beta: float, max_position_embeddings: int):
+    if isinstance(offset, mx.array) and offset.ndim > 0:
+        offset = offset[:, None]
+
+    scaling = 1 + beta * mx.log(
+        1 + mx.floor((mx.arange(size) + offset) / max_position_embeddings)
+    )
+    if scaling.ndim == 2:
+        return scaling[:, None, :, None]
+    else:
+        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(swiglu(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(PipelineMixin, 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 pipeline(self, group):
+        super().pipeline(group)
+        self.fa_idx = None
+        self.swa_idx = None
+        for e, l in enumerate(self.pipeline_layers):
+            if self.swa_idx is None and l.use_sliding:
+                self.swa_idx = e
+            elif self.fa_idx is None and not l.use_sliding:
+                self.fa_idx = e
+            if self.fa_idx is not None and self.swa_idx is not None:
+                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)
+
+        pipeline_rank = self.pipeline_rank
+        pipeline_size = self.pipeline_size
+
+        if cache is None:
+            cache = [None] * len(self.pipeline_layers)
+            offset = 0
+        else:
+            offset = cache[0].offset
+
+        swa_mask = fa_mask = None
+        if self.fa_idx is not None:
+            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(
+            inputs.shape[1],
+            offset,
+            self.args.rope_parameters["llama_4_scaling_beta"],
+            self.args.rope_parameters["original_max_position_embeddings"],
+        ).astype(h.dtype)
+
+        # 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):
+            mask = swa_mask if l.use_sliding else fa_mask
+            h = l(h, attn_scale, 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
+        if pipeline_size > 1:
+            h = mx.distributed.all_gather(h)[: h.shape[0]]
+
+        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
+
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        N = group.size()
+        for layer in self.model.layers:
+            # Shard the self attention
+            layer.self_attn.q_proj = shard_linear(
+                layer.self_attn.q_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.k_proj = shard_linear(
+                layer.self_attn.k_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.v_proj = shard_linear(
+                layer.self_attn.v_proj, "all-to-sharded", group=group
+            )
+            layer.self_attn.o_proj = shard_linear(
+                layer.self_attn.o_proj, "sharded-to-all", group=group
+            )
+            layer.self_attn.n_heads //= N
+            layer.self_attn.n_kv_heads //= N
+
+            # Shard the MLP
+            layer.mlp.gate_proj = shard_linear(
+                layer.mlp.gate_proj, "all-to-sharded", group=group
+            )
+            layer.mlp.down_proj = shard_linear(
+                layer.mlp.down_proj, "sharded-to-all", group=group
+            )
+            layer.mlp.up_proj = shard_linear(
+                layer.mlp.up_proj, "all-to-sharded", group=group
+            )
+
+    @property
+    def layers(self):
+        return self.model.pipeline_layers
+
+    def make_cache(self):
+        return [
+            (
+                RotatingKVCache(max_size=self.model.sliding_window)
+                if layer.use_sliding
+                else KVCache()
+            )
+            for layer in self.layers
+        ]

mlx_lm/models/mistral3.py

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

mlx_lm/models/nanochat.py

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

mlx_lm/models/nemotron-nas.py

@@ -7,6 +7,7 @@ import mlx.core as mx
 import mlx.nn as nn
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import KVCache
 from .rope_utils import initialize_rope
 
 
@@ -329,6 +330,9 @@ class NemotronNASModel(nn.Module):
             for i in range(args.num_hidden_layers)
         ]
         self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.num_attn_layers = sum(
+            1 for layer in self.layers if layer.self_attn is not None
+        )
 
     def __call__(
         self,
@@ -338,11 +342,17 @@ class NemotronNASModel(nn.Module):
         h = self.embed_tokens(inputs)
 
         if cache is None:
-            cache = [None] * len(self.layers)
+            cache = [None] * self.num_attn_layers
 
         mask = create_attention_mask(h, cache[0])
 
-        for layer, c in zip(self.layers, cache):
+        cache_idx = 0
+        for layer in self.layers:
+            if layer.self_attn is not None:
+                c = cache[cache_idx]
+                cache_idx += 1
+            else:
+                c = None
             h = layer(h, mask, cache=c)
 
         return self.norm(h)
@@ -380,3 +390,6 @@ class Model(nn.Module):
     @property
     def layers(self):
         return self.model.layers
+
+    def make_cache(self):
+        return [KVCache() for layer in self.layers if layer.self_attn is not None]

mlx_lm/models/nemotron_h.py

@@ -7,6 +7,7 @@ from typing import Any, List, Optional, Tuple
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import (
     BaseModelArgs,
     create_attention_mask,
@@ -15,6 +16,7 @@ from .base import (
 )
 from .cache import KVCache, MambaCache
 from .ssm import ssm_update
+from .switch_layers import SwitchMLP
 
 
 @dataclass()
@@ -37,24 +39,34 @@ class ModelArgs(BaseModelArgs):
     time_step_limit: Tuple[float, float]
     mlp_bias: bool
     layer_norm_epsilon: float
-    rms_norm_eps: float
     use_bias: bool
     use_conv_bias: bool
-    residual_in_fp32: bool
     hybrid_override_pattern: List[str]
     head_dim: Optional[int] = None
+    moe_intermediate_size: Optional[int] = None
+    moe_shared_expert_intermediate_size: Optional[int] = None
+    n_group: Optional[int] = None
+    n_routed_experts: Optional[int] = None
+    n_shared_experts: Optional[int] = None
+    topk_group: Optional[int] = None
+    num_experts_per_tok: Optional[int] = None
+    norm_topk_prob: Optional[bool] = None
+    routed_scaling_factor: Optional[float] = None
 
 
 class MambaRMSNormGated(nn.Module):
-    def __init__(self, hidden_size: int, eps: float = 1e-6):
+    def __init__(self, hidden_size: int, eps: float, group_size: int):
         super().__init__()
         self.eps = eps
         self.weight = mx.ones(hidden_size)
+        self.group_size = group_size
 
-    def __call__(self, hidden_states: mx.array, gate: mx.array = None) -> mx.array:
+    def __call__(self, x: 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)
+            x = swiglu(gate, x)
+        x = mx.unflatten(x, axis=-1, shape=(-1, self.group_size))
+        x = mx.fast.rms_norm(x, weight=None, eps=self.eps)
+        return self.weight * x.flatten(-2)
 
 
 class NemotronHMamba2Mixer(nn.Module):
@@ -90,16 +102,25 @@ class NemotronHMamba2Mixer(nn.Module):
         self.A_log = mx.log(mx.arange(1, self.num_heads + 1, dtype=mx.float32))
         self.D = mx.ones(self.num_heads)
 
+        group_size = self.intermediate_size // self.n_groups
         self.norm = MambaRMSNormGated(
-            self.intermediate_size, eps=args.layer_norm_epsilon
+            self.intermediate_size,
+            eps=args.layer_norm_epsilon,
+            group_size=group_size,
         )
         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
+    def _conv(
+        self,
+        conv_input: mx.array,
+        cache: Optional[MambaCache],
+        mask: Optional[mx.array],
     ) -> mx.array:
+        if mask is not None:
+            conv_input = mx.where(mask[..., None], conv_input, 0)
+
         if cache is not None:
             if cache[0] is None:
                 conv_state = mx.zeros(
@@ -109,11 +130,19 @@ class NemotronHMamba2Mixer(nn.Module):
             else:
                 conv_state = cache[0]
             padded_input = mx.concatenate([conv_state, conv_input], axis=1)
-            cache[0] = padded_input[:, -(self.conv_kernel_size - 1) :, :]
+            n_keep = self.conv_kernel_size - 1
+            if cache.lengths is not None:
+                t = padded_input.shape[1]
+                ends = mx.clip(cache.lengths, 0, t - n_keep)
+                positions = (ends[:, None] + mx.arange(n_keep))[..., None]
+                cache[0] = mx.take_along_axis(padded_input, positions, axis=1)
+            else:
+                cache[0] = padded_input[:, -n_keep:, :]
         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)
 
@@ -123,8 +152,8 @@ class NemotronHMamba2Mixer(nn.Module):
         B: mx.array,
         C: mx.array,
         dt: mx.array,
-        state: Optional[mx.array],
-        mask: Optional[mx.array] = None,
+        cache: Optional[MambaCache],
+        mask: Optional[mx.array],
     ) -> mx.array:
         batch_size, seq_len, _ = hidden_states.shape
 
@@ -133,21 +162,28 @@ class NemotronHMamba2Mixer(nn.Module):
         )
         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)
+        if cache:
+            state = cache[1]
+            lengths = cache.lengths
+        else:
+            state, lengths = None, None
 
         y, state = ssm_update(
             hidden_states,
             self.A_log,
             B,
             C,
-            self.D,
+            self.D.astype(hidden_states.dtype),
             dt,
             self.dt_bias,
             state,
             self.time_step_limit,
             mask,
         )
+        if cache:
+            cache[1] = state
 
-        return y.reshape(batch_size, seq_len, self.intermediate_size), state
+        return y.reshape(batch_size, seq_len, self.intermediate_size)
 
     def __call__(
         self,
@@ -163,11 +199,7 @@ class NemotronHMamba2Mixer(nn.Module):
             [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)
-
+        conv_output = self._conv(conv_input, cache, mask)
         hidden_states_ssm, B, C = mx.split(
             conv_output,
             [
@@ -176,10 +208,9 @@ class NemotronHMamba2Mixer(nn.Module):
             ],
             axis=-1,
         )
-        state = cache[1] if cache else None
-        y, state = self._ssm(hidden_states_ssm, B, C, dt, state, mask)
+        y = self._ssm(hidden_states_ssm, B, C, dt, cache, mask)
         if cache:
-            cache[1] = state
+            cache.advance(y.shape[1])
         y = self.norm(y, gate)
         return self.out_proj(y)
 
@@ -245,24 +276,113 @@ class NemotronHAttention(nn.Module):
 
 
 class NemotronHMLP(nn.Module):
-    def __init__(self, args: ModelArgs):
+    def __init__(self, args: ModelArgs, intermediate_size=None):
         super().__init__()
+        intermediate_size = intermediate_size or args.intermediate_size
+
         self.up_proj = nn.Linear(
-            args.hidden_size, args.intermediate_size, bias=args.mlp_bias
+            args.hidden_size, intermediate_size, bias=args.mlp_bias
         )
         self.down_proj = nn.Linear(
-            args.intermediate_size, args.hidden_size, bias=args.mlp_bias
+            intermediate_size, args.hidden_size, bias=args.mlp_bias
         )
 
     def __call__(self, x):
         return self.down_proj(nn.relu2(self.up_proj(x)))
 
 
+@mx.compile
+def group_expert_select(
+    gates,
+    e_score_correction_bias,
+    top_k,
+    n_group,
+    topk_group,
+    routed_scaling_factor,
+    norm_topk_prob,
+):
+
+    orig_scores = scores = mx.sigmoid(gates.astype(mx.float32))
+    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 + 1e-20)
+    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,))
+
+    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 NemotronHMoE(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 = SwitchMLP(
+            config.hidden_size,
+            config.moe_intermediate_size,
+            config.n_routed_experts,
+            activation=nn.ReLU2(),
+        )
+
+        self.gate = MoEGate(config)
+        if config.n_shared_experts is not None:
+            intermediate_size = config.moe_shared_expert_intermediate_size
+            self.shared_experts = NemotronHMLP(
+                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 NemotronHBlock(nn.Module):
     def __init__(self, args: ModelArgs, block_type: str):
         super().__init__()
-        self.residual_in_fp32 = args.residual_in_fp32
-        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon)
 
         self.block_type = block_type
 
@@ -272,6 +392,8 @@ class NemotronHBlock(nn.Module):
             self.mixer = NemotronHAttention(args)
         elif self.block_type == "-":
             self.mixer = NemotronHMLP(args)
+        elif self.block_type == "E":
+            self.mixer = NemotronHMoE(args)
 
     def __call__(
         self,
@@ -296,7 +418,7 @@ class NemotronHModel(nn.Module):
             NemotronHBlock(args, block_type)
             for block_type in args.hybrid_override_pattern
         ]
-        self.norm_f = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.norm_f = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon)
         self.fa_idx = 0
         self.ssm_idx = 0
         for b in args.hybrid_override_pattern:
@@ -372,4 +494,23 @@ class Model(nn.Module):
         for k, v in weights.items():
             if "conv1d.weight" in k and v.shape[-1] != 1:
                 weights[k] = v.moveaxis(2, 1)
+
+        # Stack experts
+        for l in range(self.args.num_hidden_layers):
+            prefix = f"backbone.layers.{l}.mixer"
+            for m, n in [("down_proj", "fc2"), ("up_proj", "fc1")]:
+                if f"{prefix}.experts.0.{m}.weight" in weights:
+                    to_join = [
+                        weights.pop(f"{prefix}.experts.{e}.{m}.weight")
+                        for e in range(self.args.n_routed_experts)
+                    ]
+                    weights[f"{prefix}.switch_mlp.{n}.weight"] = mx.stack(to_join)
+
         return weights
+
+    @property
+    def cast_predicate(self):
+        def predicate(k):
+            return "e_score_correction_bias" not in k and "A_log" not in k
+
+        return predicate

mlx_lm/models/olmo.py

@@ -7,6 +7,7 @@ from typing import Any, Optional
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask
 
 try:
@@ -105,7 +106,7 @@ class TransformerBlock(nn.Module):
 
         x1, x2 = mx.split(self.ff_proj(self.ff_norm(h)), 2, axis=-1)
 
-        out = h + self.ff_out(nn.silu(x2) * x1)
+        out = h + self.ff_out(swiglu(x2, x1))
         return out
 
 

mlx_lm/models/olmo2.py

@@ -6,6 +6,7 @@ from typing import Any, Dict, Optional, Union
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .rope_utils import initialize_rope
 
@@ -115,7 +116,7 @@ class MLP(nn.Module):
         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))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class TransformerBlock(nn.Module):

mlx_lm/models/olmo3.py

@@ -0,0 +1,236 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .activations import swiglu
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import KVCache, RotatingKVCache
+from .rope_utils import initialize_rope
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    num_hidden_layers: int
+    intermediate_size: int
+    num_attention_heads: int
+    rms_norm_eps: float
+    vocab_size: int
+    max_position_embeddings: int
+    sliding_window: int
+    rope_theta: float
+    attention_bias: bool = False
+    layer_types: Optional[List[str]] = None
+    num_key_value_heads: Optional[int] = None
+    head_dim: Optional[int] = None
+    rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = False
+
+    def __post_init__(self):
+        if self.num_key_value_heads is None:
+            self.num_key_value_heads = self.num_attention_heads
+        if self.layer_types is None:
+            self.layer_types = [
+                "full_attention" if (i + 1) % 4 == 0 else "sliding_attention"
+                for i in range(self.num_hidden_layers)
+            ]
+
+
+class Olmo3Attention(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.num_key_value_heads = args.num_key_value_heads
+        self.layer_idx = layer_idx
+
+        self.head_dim = args.head_dim or args.hidden_size // args.num_attention_heads
+        self.scale = self.head_dim**-0.5
+
+        self.q_proj = nn.Linear(
+            args.hidden_size,
+            args.num_attention_heads * self.head_dim,
+            bias=args.attention_bias,
+        )
+        self.k_proj = nn.Linear(
+            args.hidden_size,
+            args.num_key_value_heads * self.head_dim,
+            bias=args.attention_bias,
+        )
+        self.v_proj = nn.Linear(
+            args.hidden_size,
+            args.num_key_value_heads * self.head_dim,
+            bias=args.attention_bias,
+        )
+        self.o_proj = nn.Linear(
+            args.num_attention_heads * self.head_dim,
+            args.hidden_size,
+            bias=args.attention_bias,
+        )
+
+        self.q_norm = nn.RMSNorm(
+            args.num_attention_heads * self.head_dim, eps=args.rms_norm_eps
+        )
+        self.k_norm = nn.RMSNorm(
+            args.num_key_value_heads * self.head_dim, eps=args.rms_norm_eps
+        )
+
+        if args.layer_types[layer_idx] != "full_attention":
+            self.rope = nn.RoPE(self.head_dim, traditional=False, base=args.rope_theta)
+        else:
+            self.rope = initialize_rope(
+                self.head_dim,
+                traditional=False,
+                base=args.rope_theta,
+                scaling_config=args.rope_scaling,
+                max_position_embeddings=args.max_position_embeddings,
+            )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, _ = x.shape
+        queries = self.q_norm(self.q_proj(x))
+        keys = self.k_norm(self.k_proj(x))
+        values = self.v_proj(x)
+
+        queries = queries.reshape(B, L, self.num_attention_heads, -1).transpose(
+            0, 2, 1, 3
+        )
+        keys = keys.reshape(B, L, self.num_key_value_heads, -1).transpose(0, 2, 1, 3)
+        values = values.reshape(B, L, self.num_key_value_heads, -1).transpose(
+            0, 2, 1, 3
+        )
+
+        if cache is not None:
+            queries = self.rope(queries, offset=cache.offset)
+            keys = self.rope(keys, offset=cache.offset)
+            keys, values = cache.update_and_fetch(keys, values)
+        else:
+            queries = self.rope(queries)
+            keys = self.rope(keys)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
+
+
+class Olmo3MLP(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.gate_proj = nn.Linear(args.hidden_size, args.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(args.intermediate_size, args.hidden_size, bias=False)
+        self.up_proj = nn.Linear(args.hidden_size, args.intermediate_size, bias=False)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
+
+
+class Olmo3DecoderLayer(nn.Module):
+    def __init__(self, args: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.num_attention_heads = args.num_attention_heads
+        self.hidden_size = args.hidden_size
+        self.self_attn = Olmo3Attention(args, layer_idx=layer_idx)
+        self.mlp = Olmo3MLP(args)
+        self.post_attention_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+        self.post_feedforward_layernorm = nn.RMSNorm(
+            args.hidden_size, eps=args.rms_norm_eps
+        )
+        self.args = args
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        r = self.post_attention_layernorm(self.self_attn(x, mask, cache))
+        h = x + r
+        r = self.post_feedforward_layernorm(self.mlp(h))
+        out = h + r
+        return out
+
+
+class Olmo3Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.sliding_window = args.sliding_window
+
+        self.embed_tokens = nn.Embedding(args.vocab_size, args.hidden_size)
+        self.layers = [
+            Olmo3DecoderLayer(args=args, layer_idx=i)
+            for i in range(args.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+
+        self.swa_idx = args.layer_types.index("sliding_attention")
+        self.ga_idx = args.layer_types.index("full_attention")
+        self.layer_types = args.layer_types
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        h = self.embed_tokens(inputs)
+
+        if cache is None:
+            cache = [None] * len(self.layers)
+
+        full_mask = create_attention_mask(h, cache[self.ga_idx])
+        sliding_window_mask = create_attention_mask(
+            h, cache[self.swa_idx], window_size=self.sliding_window
+        )
+
+        for layer, c, layer_type in zip(self.layers, cache, self.layer_types):
+            mask = full_mask if layer_type == "full_attention" else sliding_window_mask
+            h = layer(h, mask, cache=c)
+
+        return self.norm(h)
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.model = Olmo3Model(args)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        out = self.model(inputs, cache)
+        if self.args.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
+        return out
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    def make_cache(self):
+        caches = []
+        for lt in self.model.layer_types:
+            if lt == "full_attention":
+                caches.append(KVCache())
+            else:
+                caches.append(RotatingKVCache(max_size=self.args.sliding_window))
+        return caches

mlx_lm/models/openelm.py

@@ -6,6 +6,7 @@ from typing import Any, Dict, List, Optional, Union
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
@@ -136,7 +137,7 @@ class MLP(nn.Module):
     def __call__(self, x) -> mx.array:
         x = self.proj_1(x)
         gate, x = mx.split(x, 2, axis=-1)
-        return self.proj_2(nn.silu(gate) * x)
+        return self.proj_2(swiglu(gate, x))
 
 
 class TransformerBlock(nn.Module):

mlx_lm/models/phi3.py

@@ -6,6 +6,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union
 import mlx.core as mx
 import mlx.nn as nn
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .rope_utils import SuScaledRoPE
 
@@ -126,7 +127,7 @@ class MLP(nn.Module):
     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)
+        return self.down_proj(swiglu(gate, x))
 
 
 class TransformerBlock(nn.Module):

mlx_lm/models/pipeline.py

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

mlx_lm/models/plamo.py

@@ -7,6 +7,7 @@ import mlx.core as mx
 import mlx.nn as nn
 import numpy as np
 
+from .activations import swiglu
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 
 
@@ -115,7 +116,7 @@ class MLP(nn.Module):
         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))  # type: ignore
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))  # type: ignore
 
 
 class PlamoDecoderLayer(nn.Module):