added tqdm to save_model

Co-authored-by: Tristan Stahnke <tristan@melchior.lan>
Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>

.github/workflows/pull_request.yml

@@ -38,4 +38,7 @@ jobs:
       - name: Run tests
         shell: bash -l {0}
         run: |
-          python -m xmlrunner discover -v tests -o test-results/
+          curl -o test_data.zip -L https://github.com/ml-explore/mlx-lm/releases/download/test_data/test_data.zip
+          unzip test_data.zip
+          METAL_DEVICE_WRAPPER_TYPE=1 METAL_DEBUG_ERROR_MODE=0 HF_HOME="." python -m xmlrunner discover -v tests -o test-results/
+          mlx.launch -n 2 tests/model_parallel_tests.py

ACKNOWLEDGMENTS.md

@@ -10,11 +10,11 @@ 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` and 
-`Mamba v2`, Z.ai & THUKEG's `GLM`, `GLM4`, Rednote `dots.llm1`, Baidu's `Ernie4.5 MoE`, 
+`Mamba v2`, Z.ai & THUKEG's `GLM`, `GLM4`, `GLM5 (GLM MoE DSA)`, 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`; 
+Alibaba Qwen's `Qwen3Next`, Tele-AI's `TeleChat3`, 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; 
@@ -26,4 +26,8 @@ Added support for the following other features:
   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`.
+ ServiceNow-AI's `Apriel 1.5`, Tencent's `Hunyuan Dense V1` and `Hunyuan MoE V1`.
+- Tarjei Mandt: Added support for the following architectures: `Step 3.5 Flash`, 
+MoonshotAI's `Kimi K2.5`, Upstage's `Solar Open`, LG AI Research's `K-Exaone MoE`, 
+Meituan's `LongCat Flash Lite` Helped add support for the following model architectures: 
+Z.ai & THUKEG's `GLM5 (GLM MoE DSA)`

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
 ```

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/LORA.md

@@ -66,9 +66,10 @@ mlx_lm.lora \
 To fine-tune the full model weights, add the `--fine-tune-type full` flag.
 Currently supported fine-tuning types are `lora` (default), `dora`, and `full`.
 
-The `--data` argument must specify a path to a `train.jsonl`, `valid.jsonl`
-when using `--train` and a path to a `test.jsonl` when using `--test`. For more
-details on the data format see the section on [Data](#Data).
+The `--data` argument must specify a path to a `train.jsonl` when using
+`--train` and a path to a `test.jsonl` when using `--test`. A `valid.jsonl` is
+optional; if provided, validation loss will be reported during training. For
+more details on the data format see the section on [Data](#Data).
 
 For example, to fine-tune a Mistral 7B you can use `--model
 mistralai/Mistral-7B-v0.1`.
@@ -184,9 +185,10 @@ Face.
 
 ### Local Datasets
 
-For fine-tuning (`--train`), the data loader expects a `train.jsonl` and a
-`valid.jsonl` to be in the data directory. For evaluation (`--test`), the data
-loader expects a `test.jsonl` in the data directory. 
+For fine-tuning (`--train`), the data loader expects a `train.jsonl` to be in
+the data directory. A `valid.jsonl` is optional; if present, validation loss
+will be reported periodically during training. For evaluation (`--test`), the
+data loader expects a `test.jsonl` in the data directory. 
 
 Currently, `*.jsonl` files support `chat`, `tools`, `completions`, and `text`
 data formats. Here are examples of these formats:

mlx_lm/SERVER.md

@@ -72,12 +72,24 @@ curl localhost:8080/v1/chat/completions \
 - `min_p`: (Optional) A float specifying the min-p sampling parameter.
   Defaults to `0.0` (disabled).
 
-- `repetition_penalty`: (Optional) Applies a penalty to repeated tokens.
-  Defaults to `1.0`.
+- `repetition_penalty`: (Optional) Applies a multiplicative penalty to repeated
+  tokens. Defaults to `0.0` (disabled).
 
 - `repetition_context_size`: (Optional) The size of the context window for
   applying repetition penalty. Defaults to `20`.
 
+- `presence_penalty`: (Optional) Applies an additive penalty to tokens
+  that appeared before. Defaults to `0.0` (disabled).
+
+- `presence_context_size`: (Optional) The size of the context window for
+  applying presence penalty. Defaults to `20`.
+
+- `frequency_penalty`: (Optional) Applies an additive penalty proportional to
+  how many times a token appeared previously. Defaults to `0.0` (disabled).
+
+- `frequency_context_size`: (Optional) The size of the context window for
+  applying frequency penalty. Defaults to `20`.
+
 - `logit_bias`: (Optional) A dictionary mapping token IDs to their bias
   values. Defaults to `None`.
 

mlx_lm/__main__.py

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

mlx_lm/_version.py

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

mlx_lm/benchmark.py

@@ -1,12 +1,13 @@
 # Copyright © 2025 Apple Inc.
 
 import argparse
+import time
 
 import mlx.core as mx
 
 from mlx_lm import batch_generate, load, stream_generate
 from mlx_lm.generate import DEFAULT_MODEL
-from mlx_lm.utils import pipeline_load
+from mlx_lm.utils import pipeline_load, sharded_load
 
 
 def setup_arg_parser():
@@ -49,6 +50,29 @@ 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",
+    )
+    parser.add_argument(
+        "--quantize-activations",
+        "-qa",
+        action="store_true",
+        help="Quantize activations using the same quantization config as the corresponding layer.",
+    )
+    parser.add_argument(
+        "--prefill-step-size",
+        type=int,
+        default=2048,
+        help="Step size for prefill processing (default: 2048)",
+    )
+    parser.add_argument(
+        "--delay",
+        type=int,
+        default=0,
+        help="Delay between each test in seconds (default: 0)",
+    )
     return parser
 
 
@@ -59,6 +83,8 @@ def main():
 
     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:
@@ -67,10 +93,15 @@ def main():
     model_path = args.model or DEFAULT_MODEL
 
     if group.size() > 1:
-        model, tokenizer, config = pipeline_load(args.model, return_config=True)
+        model, tokenizer, config = sharded_load(
+            model_path, pipeline_group, tensor_group, return_config=True
+        )
     else:
         model, tokenizer, config = load(
-            args.model, return_config=True, tokenizer_config={"trust_remote_code": True}
+            model_path,
+            return_config=True,
+            tokenizer_config={"trust_remote_code": True},
+            model_config={"quantize_activations": args.quantize_activations},
         )
 
     # Empty to avoid early stopping
@@ -85,14 +116,22 @@ def main():
 
     def single_bench():
         for response in stream_generate(
-            model, tokenizer, prompt, max_tokens=generation_tokens
+            model,
+            tokenizer,
+            prompt,
+            max_tokens=generation_tokens,
+            prefill_step_size=args.prefill_step_size,
         ):
             pass
         return response
 
     def batch_bench():
         return batch_generate(
-            model, tokenizer, prompts, max_tokens=generation_tokens
+            model,
+            tokenizer,
+            prompts,
+            max_tokens=generation_tokens,
+            prefill_step_size=args.prefill_step_size,
         ).stats
 
     if batch_size == 1:
@@ -107,10 +146,15 @@ def main():
     rprint(f"Timing with {prompt_tokens=}, {generation_tokens=}, {batch_size=}.")
     responses = []
     for i in range(args.num_trials):
+        if args.delay > 0:
+            time.sleep(args.delay)
+        tic = time.perf_counter()
         response = _bench()
+        toc = time.perf_counter()
         responses.append(response)
         results = [(k, getattr(response, k)) for k in report_keys]
         results = [f"{k}={v:.3f}" for k, v in results]
+        results.append(f"total_time={toc - tic:.3f}")
         rprint(f"Trial {i+1}:  " + ", ".join(results))
 
     def avg(k):

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,13 +7,13 @@ 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
 DEFAULT_XTC_PROBABILITY = 0.0
 DEFAULT_XTC_THRESHOLD = 0.0
-DEFAULT_SEED = None
+DEFAULT_SEED = 0
 DEFAULT_MAX_TOKENS = 256
 DEFAULT_MODEL = "mlx-community/Llama-3.2-3B-Instruct-4bit"
 
@@ -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,41 @@ def main():
     parser = setup_arg_parser()
     args = parser.parse_args()
 
-    if args.seed is not None:
-        mx.random.seed(args.seed)
+    group = mx.distributed.init()
+    rank = group.rank()
+    pipeline_group = group if args.pipeline else None
+    tensor_group = group if not args.pipeline else None
 
-    model, tokenizer = load(
-        args.model,
-        adapter_path=args.adapter_path,
-        tokenizer_config={
-            "trust_remote_code": True if args.trust_remote_code else None
-        },
-    )
+    def rprint(*args, **kwargs):
+        if rank == 0:
+            print(*args, **kwargs)
+
+    mx.random.seed(args.seed)
+
+    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 +138,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 +158,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,345 @@
+# Copyright © 2025 Apple Inc.
+
+import copy
+import json
+import re
+from inspect import isfunction
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+from transformers.utils.chat_template_utils import get_json_schema
+
+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):
+    def normalize_tool(tool):
+        if isfunction(tool):
+            return get_json_schema(tool)
+        return tool["function"]
+
+    return [normalize_tool(tool) 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,
+    tools: Any = None,
+) -> 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 = tools or msg.get("tools")
+    response_format = msg.get("response_format")
+    tool_calls = msg.get("tool_calls")
+    reasoning_content = msg.get("reasoning_content")
+
+    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_from_openai_format(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_from_openai_format(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,
+    tools: Any = None,
+) -> 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,
+            tools=tools,
+        )
+
+    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/cli.py

@@ -0,0 +1,53 @@
+# Copyright © 2025 Apple Inc.
+
+import importlib
+import sys
+
+
+def main():
+    subcommands = (
+        "benchmark",
+        "cache_prompt",
+        "chat",
+        "convert",
+        "evaluate",
+        "fuse",
+        "generate",
+        "lora",
+        "manage",
+        "perplexity",
+        "awq",
+        "dwq",
+        "dynamic_quant",
+        "gptq",
+        "server",
+        "upload",
+        "share",
+    )
+    subpackages = {
+        "awq": "quant",
+        "dwq": "quant",
+        "dynamic_quant": "quant",
+        "gptq": "quant",
+    }
+    if len(sys.argv) < 2:
+        raise ValueError(f"CLI requires a subcommand in {subcommands}")
+    subcommand = sys.argv.pop(1)
+    if subcommand in subcommands:
+        if subpackage := subpackages.get(subcommand):
+            subcommand = f"{subpackage}.{subcommand}"
+        submodule = importlib.import_module(f"mlx_lm.{subcommand}")
+        submodule.main()
+    elif subcommand == "--version":
+        from mlx_lm import __version__
+
+        print(__version__)
+    elif subcommand in ("-h", "--help"):
+        print(f"The supported subcommands are {subcommands}")
+        print()
+        print(
+            "For help on an individual subcommand, pass --help "
+            "to the subcommand. For example: mlx_lm.generate --help"
+        )
+    else:
+        raise ValueError(f"CLI requires a subcommand in {subcommands}")

mlx_lm/convert.py

@@ -20,6 +20,7 @@ from .utils import (
 def mixed_quant_predicate_builder(
     recipe: str, model: nn.Module, group_size: int = 64
 ) -> Callable[[str, nn.Module, dict], Union[bool, dict]]:
+    mode = "affine"
     high_bits = 6
 
     if recipe == "mixed_2_6":
@@ -65,13 +66,13 @@ def mixed_quant_predicate_builder(
         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}
+            return {"group_size": group_size, "bits": high_bits, "mode": mode}
         if "down_proj" in path and use_more_bits:
-            return {"group_size": group_size, "bits": high_bits}
+            return {"group_size": group_size, "bits": high_bits, "mode": mode}
         if "lm_head" in path:
-            return {"group_size": group_size, "bits": high_bits}
+            return {"group_size": group_size, "bits": high_bits, "mode": mode}
 
-        return {"group_size": group_size, "bits": low_bits}
+        return {"group_size": group_size, "bits": low_bits, "mode": mode}
 
     return mixed_quant_predicate
 
@@ -85,8 +86,8 @@ def convert(
     hf_path: str,
     mlx_path: str = "mlx_model",
     quantize: bool = False,
-    q_group_size: int = 64,
-    q_bits: int = 4,
+    q_group_size: Optional[int] = None,
+    q_bits: Optional[int] = None,
     q_mode: str = "affine",
     dtype: Optional[str] = None,
     upload_repo: str = None,
@@ -117,12 +118,18 @@ def convert(
     )
 
     if isinstance(quant_predicate, str):
+        if q_mode != "affine":
+            raise ValueError(f"Quant predicates only support 'affine' quantization.")
         quant_predicate = mixed_quant_predicate_builder(
-            quant_predicate, model, q_group_size
+            quant_predicate,
+            model,
+            q_group_size,
         )
 
     if dtype is None:
         dtype = config.get("torch_dtype", None)
+    if dtype is None and (text_config := config.get("text_config", None)):
+        dtype = text_config.get("dtype", None)
     if dtype in MODEL_CONVERSION_DTYPES:
         print("[INFO] Using dtype:", dtype)
         dtype = getattr(mx, dtype)
@@ -179,7 +186,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."
     )
@@ -187,17 +199,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

@@ -20,7 +20,6 @@ import mlx.nn as nn
 import numpy as np
 from lm_eval.api.model import LM
 from lm_eval.api.registry import register_model
-from lm_eval.models import huggingface
 from tqdm import tqdm
 
 from .generate import batch_generate
@@ -72,13 +71,13 @@ def chat_template_fn(**extra_kwargs):
 @register_model("mlxlm")
 class MLXLM(LM):
 
-    tokenizer_name = huggingface.HFLM.tokenizer_name
     apply_chat_template = chat_template_fn()
 
     def __init__(
         self,
         path_or_hf_repo: str,
         max_tokens: Optional[int] = None,
+        batch_size: int = 8,
         use_chat_template: Optional[bool] = None,
         trust_remote_code: bool = False,
         sampler: Optional[Callable[[mx.array], mx.array]] = None,
@@ -89,7 +88,7 @@ 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
@@ -146,6 +145,10 @@ class MLXLM(LM):
             for t in texts
         ]
 
+    @property
+    def tokenizer_name(self) -> str:
+        return self.tokenizer.name_or_path.replace("/", "__")
+
     def loglikelihood(self, requests) -> list[tuple[float, bool]]:
         """Compute log-likelihood of generating a continuation from a context.
         Downstream tasks should attempt to use loglikelihood instead of other
@@ -476,6 +479,7 @@ def main():
     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,

mlx_lm/examples/batch_generate_response.py

@@ -27,7 +27,7 @@ prompts = [
 
 # Set `verbose=True` to see generation statistics
 result = batch_generate(
-    model, tokenizer, prompts, verbose=False, return_prompt_caches=True
+    model, tokenizer, prompts, verbose=False, return_prompt_caches=True, max_tokens=2048
 )
 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/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/sharded_generate.py

@@ -1,19 +1,20 @@
-# Copyright © 2024 Apple Inc.
+# Copyright © 2025 Apple Inc.
 
 """
 Run with:
 
 ```
 mlx.launch \
- --hostfile /path/to/hosts.json \
- /path/to/pipeline_generate.py \
- --prompt "hello world"
+    --backend jaccl \
+    --env MLX_METAL_FAST_SYNCH=1 \
+    --hostfile /path/to/hosts.json \
+    /path/to/sharded_generate.py \
+    --prompt 'Hello world'
 ```
 
-Make sure you can run MLX over MPI on two hosts. For more information see the
-documentation:
+For more information on running distributed programs with MLX see the documentation:
 
-https://ml-explore.github.io/mlx/build/html/usage/distributed.html).
+https://ml-explore.github.io/mlx/build/html/usage/distributed.html .
 """
 
 import argparse
@@ -21,13 +22,13 @@ import argparse
 import mlx.core as mx
 
 from mlx_lm import stream_generate
-from mlx_lm.utils import pipeline_load
+from mlx_lm.utils import sharded_load
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="LLM pipelined inference example")
+    parser = argparse.ArgumentParser(description="LLM distributed inference example")
     parser.add_argument(
         "--model",
-        default="mlx-community/DeepSeek-R1-3bit",
+        default="mlx-community/Llama-3.3-70B-Instruct-4bit",
         help="HF repo or path to local model.",
     )
     parser.add_argument(
@@ -43,19 +44,29 @@ if __name__ == "__main__":
         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 = pipeline_load(args.model)
+    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)
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+    )
 
     for response in stream_generate(
         model, tokenizer, prompt, max_tokens=args.max_tokens

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

@@ -76,8 +76,9 @@ def main() -> None:
 
     if args.dequantize:
         print("Dequantizing model")
-        model = dequantize(model)
+        model = dequantize_model(model)
         config.pop("quantization", None)
+        config.pop("quantization_config", None)
 
     save_path = Path(args.save_path)
     save(

mlx_lm/lora.py

@@ -21,7 +21,7 @@ from .tuner.utils import (
     load_adapters,
     print_trainable_parameters,
 )
-from .utils import load, save_config
+from .utils import _parse_size, load, save_config
 
 yaml_loader = yaml.SafeLoader
 yaml_loader.add_implicit_resolver(
@@ -69,6 +69,7 @@ CONFIG_DEFAULTS = {
     "config": None,
     "grad_checkpoint": False,
     "grad_accumulation_steps": 1,
+    "clear_cache_threshold": 0,
     "lr_schedule": None,
     "lora_parameters": {"rank": 8, "dropout": 0.0, "scale": 20.0},
     "mask_prompt": False,
@@ -190,6 +191,12 @@ def build_parser():
         help="Use gradient checkpointing to reduce memory use.",
         default=None,
     )
+    parser.add_argument(
+        "--clear-cache-threshold",
+        type=_parse_size,
+        default=0,
+        help="Clear the allocator cache between steps if it grows too large.",
+    )
     parser.add_argument(
         "--report-to",
         type=str,

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,43 @@
+# 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
+
+
+@partial(mx.compile, shapeless=True)
+def xielu(x, alpha_p, alpha_n, beta, eps):
+    alpha_p = nn.softplus(alpha_p)
+    alpha_n = beta + nn.softplus(alpha_n)
+    return mx.where(
+        x > 0,
+        alpha_p * mx.square(x) + beta * x,
+        (mx.expm1(mx.minimum(x, eps)) - x) * alpha_n + beta * x,
+    )
+
+
+class XieLU(nn.Module):
+    def __init__(
+        self,
+        alpha_p_init=0.8,
+        alpha_n_init=0.8,
+        beta=0.5,
+        eps=-1e-6,
+    ):
+        super().__init__()
+        alpha_p_tensor = mx.array(alpha_p_init)
+        alpha_n_tensor = mx.array(alpha_n_init - beta)
+        self.alpha_p = mx.log(mx.exp(alpha_p_tensor) - 1)
+        self.alpha_n = mx.log(mx.exp(alpha_n_tensor) - 1)
+
+        self.beta = mx.array(beta)
+        self.eps = mx.array(eps)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return xielu(x, self.alpha_p, self.alpha_n, self.beta, self.eps)

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
@@ -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

@@ -7,6 +7,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
 from .cache import KVCache, RotatingKVCache
 from .rope_utils import initialize_rope
@@ -149,7 +150,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 MoERouter(nn.Module):

mlx_lm/models/apertus.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 XieLU
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .rope_utils import initialize_rope
 
@@ -32,38 +33,6 @@ class ModelArgs(BaseModelArgs):
     rope_scaling: Optional[Dict[str, Union[float, str]]] = None
 
 
-@partial(mx.compile, shapeless=True)
-def xielu(x, alpha_p, alpha_n, beta, eps):
-    alpha_p = nn.softplus(alpha_p)
-    alpha_n = beta + nn.softplus(alpha_n)
-    return mx.where(
-        x > 0,
-        alpha_p * mx.square(x) + beta * x,
-        (mx.expm1(mx.minimum(x, eps)) - x) * alpha_n + beta * x,
-    )
-
-
-class XieLU(nn.Module):
-    def __init__(
-        self,
-        alpha_p_init=0.8,
-        alpha_n_init=0.8,
-        beta=0.5,
-        eps=-1e-6,
-    ):
-        super().__init__()
-        alpha_p_tensor = mx.array(alpha_p_init)
-        alpha_n_tensor = mx.array(alpha_n_init - beta)
-        self.alpha_p = mx.log(mx.exp(alpha_p_tensor) - 1)
-        self.alpha_n = mx.log(mx.exp(alpha_n_tensor) - 1)
-
-        self.beta = mx.array(beta)
-        self.eps = mx.array(eps)
-
-    def __call__(self, x: mx.array) -> mx.array:
-        return xielu(x, self.alpha_p, self.alpha_n, self.beta, self.eps)
-
-
 class ApertusMLP(nn.Module):
     def __init__(self, args: ModelArgs):
         super().__init__()
@@ -198,7 +167,8 @@ class Model(nn.Module):
         self.args = args
         self.model_type = args.model_type
         self.model = ApertusModel(args)
-        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+        if not args.tie_word_embeddings:
+            self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
 
     def __call__(
         self,
@@ -206,12 +176,18 @@ class Model(nn.Module):
         cache: Optional[Any] = None,
     ) -> mx.array:
         out = self.model(inputs, cache)
-        return self.lm_head(out)
+        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):
         for k, v in weights.items():
             if k.endswith("alpha_p") or k.endswith("alpha_n"):
                 weights[k] = v.squeeze()
+        if self.args.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
         return weights
 
     @property

mlx_lm/models/baichuan_m1.py

@@ -6,8 +6,9 @@ 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
+from .cache import ArraysCache, CacheList, KVCache, RotatingKVCache
 
 
 @dataclass
@@ -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):
@@ -222,7 +223,7 @@ class Model(nn.Module):
         caches = []
         for i, layer in enumerate(self.model.layers):
             is_swa = i in self.config.sliding_window_layers
-            conv_cache = MambaCache()
+            conv_cache = ArraysCache(size=2)
             if is_swa:
                 kv_cache = RotatingKVCache(max_size=self.config.sliding_window)
             else:

mlx_lm/models/bailing_moe.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
@@ -49,11 +50,6 @@ class ModelArgs(BaseModelArgs):
     moe_router_enable_shared_expert: bool = True
 
 
-@partial(mx.compile, shapeless=True)
-def swiglu(gate, up):
-    return nn.silu(gate) * up
-
-
 @partial(mx.compile, shapeless=True)
 def aggregate_expert_outputs(expert_outputs, scores):
     return (

mlx_lm/models/bailing_moe_linear.py

@@ -7,6 +7,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,
@@ -130,7 +131,7 @@ class MLP(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 Attention(nn.Module):

mlx_lm/models/cache.py

@@ -1,11 +1,13 @@
 # Copyright © 2023-2024 Apple Inc.
 
 import copy
+from collections import deque
+from dataclasses import dataclass
 from typing import Any, Dict, List, Optional
 
 import mlx.core as mx
 import mlx.nn as nn
-from mlx.utils import tree_flatten, tree_map, tree_unflatten
+from mlx.utils import tree_flatten, tree_map, tree_reduce, tree_unflatten
 
 from .base import create_causal_mask
 
@@ -109,16 +111,14 @@ def trim_prompt_cache(cache: List[Any], num_tokens: int) -> List[Any]:
     return [c.trim(num_tokens) for c in cache][0]
 
 
-def cache_length(cache: List[Any]):
-    return max(len(c) for c in cache)
-
-
 def create_attention_mask(
     N: int, offset: int, return_array: bool, window_size: Optional[int]
 ):
-    if N == 1:
+    if window_size is not None:
+        return create_causal_mask(N, offset, window_size=window_size)
+    elif N == 1:
         return None
-    if return_array:
+    elif return_array:
         return create_causal_mask(N, offset, window_size=window_size)
     else:
         return "causal"
@@ -146,23 +146,25 @@ class _BaseCache:
     def is_trimmable(self):
         return False
 
-    def __len__(self):
-        """The length of a cache is meant to represent the number of elements
-        that we need to process in the attention. For instance for KVCache it
-        is the size of the state, for RotatingKVCache it would be up to
-        max_size etc."""
+    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 __bool__(self):
-        """When an object defines __len__ then python defines the bool operator
-        as len(obj) != 0. This, for instance, doesn't allow us to write
+    @property
+    def nbytes(self):
+        """Return the size of this cache in bytes"""
+        raise NotImplementedError("Cache sub-class must implement nbytes")
 
-            cache = cache or make_cache()
-
-        which is why we are overriding that behaviour with a constant bool
-        operator return True.
+    def empty(self):
         """
-        return True
+        Return if the cache is empty or not.
+        """
+        raise NotImplementedError("Cache sub-class must implement this.")
 
     @classmethod
     def from_state(cls, state, meta_state):
@@ -217,6 +219,15 @@ 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
+
+    @property
+    def nbytes(self):
+        if self.keys is None:
+            return 0
+        return self.keys.nbytes + self.values.nbytes
+
 
 class QuantizedKVCache(_BaseCache):
     step = 256
@@ -303,6 +314,13 @@ 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
+
+    @property
+    def nbytes(self):
+        return tree_reduce(lambda a, x: a + x.nbytes, (self.keys, self.values), 0)
+
 
 class KVCache(_BaseCache):
     step = 256
@@ -336,7 +354,7 @@ class KVCache(_BaseCache):
         self.values[..., prev : self.offset, :] = values
         return self.keys[..., : self.offset, :], self.values[..., : self.offset, :]
 
-    def __len__(self):
+    def size(self):
         return self.offset
 
     @property
@@ -375,6 +393,19 @@ 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
+
+    @property
+    def nbytes(self):
+        if self.keys is None:
+            return 0
+        return self.keys.nbytes + self.values.nbytes
+
 
 class RotatingKVCache(_BaseCache):
     step = 256
@@ -483,7 +514,7 @@ class RotatingKVCache(_BaseCache):
             return self._update_in_place(keys, values)
         return self._update_concat(keys, values)
 
-    def __len__(self):
+    def size(self):
         return min(self.offset, self.max_size)
 
     @property
@@ -546,11 +577,43 @@ class RotatingKVCache(_BaseCache):
                 mask = mx.roll(mask, shift=idx + 1)
                 return mask
 
+    @classmethod
+    def merge(_, caches):
+        return BatchRotatingKVCache.merge(caches)
+
+    def empty(self):
+        return self.keys is None
+
+    @property
+    def nbytes(self):
+        if self.keys is None:
+            return 0
+        return self.keys.nbytes + self.values.nbytes
+
 
 class ArraysCache(_BaseCache):
+    def __new__(cls, *args, **kwargs):
+        instance = super().__new__(cls)
+        instance.left_padding = None
+        instance.lengths = None
+        return instance
+
     def __init__(self, size, left_padding: Optional[List[int]] = None):
         self.cache = [None] * size
-        self.left_padding = mx.array(left_padding) if left_padding else None
+        if left_padding:
+            self.left_padding = mx.array(left_padding)
+
+    @property
+    def batch_size(self):
+        for c in self.cache:
+            if c is not None:
+                return c.shape[0]
+        if self.left_padding is not None:
+            return self.left_padding.size
+        elif self.lengths is not None:
+            return self.lengths.size
+        else:
+            return 1
 
     def __setitem__(self, idx, value):
         self.cache[idx] = value
@@ -570,31 +633,108 @@ 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
+        self.cache = [c[batch_indices] if c is not None else None for c in self.cache]
+        if self.left_padding is not None:
+            self.left_padding = self.left_padding[batch_indices]
+        if self.lengths is not None:
+            self.lengths = self.lengths[batch_indices]
 
     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)]
+
+        a_batch = self.batch_size
+        b_batch = other.batch_size
+
+        def cat(a, b):
+            shape = dtype = None
+            if a is not None:
+                shape = a.shape
+                dtype = a.dtype
+            if b is not None:
+                shape = b.shape
+                dtype = b.dtype
+
+            if shape is None:
+                return None
+
+            if a is None:
+                a = mx.zeros((a_batch,) + shape[1:], dtype=dtype)
+            if b is None:
+                b = mx.zeros((b_batch,) + shape[1:], dtype=dtype)
+
+            return mx.concatenate([a, b])
+
+        self.cache = [cat(c, o) for c, o in zip(self.cache, other.cache)]
+        self.left_padding = cat(self.left_padding, other.left_padding)
+        self.lengths = cat(self.lengths, other.lengths)
+
+    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)
 
-class MambaCache(ArraysCache):
-    def __init__(self, left_padding: Optional[List[int]] = None):
-        super().__init__(size=2, left_padding=left_padding)
+        # All caches are empty so return early
+        if all(c.empty() for c in caches):
+            cache.left_padding = mx.array([0] * B)
+            return cache
+
+        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
+
+    @property
+    def nbytes(self):
+        return sum(c.nbytes for c in self.cache if c is not None)
 
 
-class ChunkedKVCache(KVCache):
+class ChunkedKVCache(_BaseCache):
+    step = 256
+
     def __init__(self, chunk_size):
-        super().__init__()
+        self.keys = None
+        self.values = None
+        self.offset = 0
         self.chunk_size = chunk_size
         self.start_position = 0
 
@@ -630,6 +770,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
@@ -643,6 +801,15 @@ 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
+
+    @property
+    def nbytes(self):
+        if self.keys is None:
+            return 0
+        return self.keys.nbytes + self.values.nbytes
+
 
 class CacheList(_BaseCache):
     def __init__(self, *caches):
@@ -661,16 +828,24 @@ class CacheList(_BaseCache):
 
     @property
     def state(self):
-        return [s for c in self.caches for s in c.state]
+        return [c.state for c in self.caches]
 
     @state.setter
     def state(self, v):
-        state_lens = [len(c.state) for c in self.caches]
-        start = 0
-        for c in self.caches:
-            l = len(c.state)
-            c.state = v[start : start + l]
-            start += l
+        for c, s in zip(self.caches, v):
+            c.state = s
+
+    @property
+    def meta_state(self):
+        return (
+            [type(c).__name__ for c in self.caches],
+            [c.meta_state for c in self.caches],
+        )
+
+    @meta_state.setter
+    def meta_state(self, v):
+        for c, m in zip(self.caches, v[1]):
+            c.meta_state = m
 
     def filter(self, batch_indices):
         """
@@ -686,6 +861,44 @@ class CacheList(_BaseCache):
         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()
+
+    @property
+    def nbytes(self):
+        return sum(c.nbytes for c in self.caches)
+
+    @classmethod
+    def from_state(cls, state, meta_state):
+        obj = cls.__new__(cls)
+        obj.caches = [
+            globals()[c].from_state(s, m) for s, c, m in zip(state, *meta_state)
+        ]
+        return obj
+
 
 def dynamic_roll(x, shifts, axis):
     n = x.shape[axis]
@@ -751,9 +964,6 @@ class BatchKVCache(_BaseCache):
         self.values[..., prev : self._idx, :] = values
         return self.keys[..., : self._idx, :], self.values[..., : self._idx, :]
 
-    def __len__(self):
-        return self._idx
-
     def prepare(self, *, left_padding=None, lengths=None, right_padding=None):
         if left_padding is not None:
             if self.keys is not None:
@@ -807,16 +1017,18 @@ class BatchKVCache(_BaseCache):
         """
         In-place filter to keep just the given indices in the cache.
         """
-        self.keys = self.keys[batch_indices]
-        self.values = self.values[batch_indices]
+        if self.keys is not None:
+            self.keys = self.keys[batch_indices]
+            self.values = self.values[batch_indices]
         self.offset = self.offset[batch_indices]
         self.left_padding = self.left_padding[batch_indices]
 
         # Shift left to reduce padding
         min_left_pad = self.left_padding.min().item()
         if min_left_pad > 0:
-            self.keys = self.keys[..., min_left_pad:, :]
-            self.values = self.values[..., min_left_pad:, :]
+            if self.keys is not None:
+                self.keys = self.keys[..., min_left_pad:, :]
+                self.values = self.values[..., min_left_pad:, :]
             self._idx -= min_left_pad
             self.left_padding -= min_left_pad
 
@@ -824,15 +1036,31 @@ class BatchKVCache(_BaseCache):
         """
         In-place extend this cache with the other cache.
         """
+        if self.keys is None and other.keys is None:
+            self.left_padding = mx.concatenate([self.left_padding, other.left_padding])
+            self.offset = mx.concatenate([self.offset, other.offset])
+            return
+
         max_idx = max(self._idx, other._idx)
-        max_size = max(self.keys.shape[2], other.keys.shape[2])
+        L1 = L2 = 0
+        if self.keys is not None:
+            B, H, L1, D = self.keys.shape
+            M = self.values.shape[3]
+        if other.keys is not None:
+            B, H, L2, D = other.keys.shape
+            M = other.values.shape[3]
+        max_size = max(L1, L2)
 
         # Pad the keys and values so they are right-justified
         # with the index and the same size
         def pad(c):
-            left = max_idx - c._idx
-            right = max_size - c.keys.shape[2] - left
             k, v = c.keys, c.values
+            if k is None:
+                Bc = c.offset.shape[0]
+                k = mx.array([]).reshape(Bc, H, 0, D)
+                v = mx.array([]).reshape(Bc, H, 0, M)
+            left = max_idx - c._idx
+            right = max_size - k.shape[2] - left
             if right < 0:
                 k = k[..., :right, :]
                 v = v[..., :right, :]
@@ -859,8 +1087,13 @@ class BatchKVCache(_BaseCache):
 
     @classmethod
     def merge(cls, caches):
-        lengths = [len(c) for c in caches]
+        lengths = [c.size() for c in caches]
         max_length = max(lengths)
+
+        # No cache has content so make an empty one
+        if max_length == 0:
+            return BatchKVCache([0] * len(caches))
+
         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)
@@ -871,6 +1104,8 @@ class BatchKVCache(_BaseCache):
         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, :]
 
@@ -882,6 +1117,18 @@ class BatchKVCache(_BaseCache):
 
         return cache
 
+    def size(self):
+        return self._idx
+
+    def empty(self):
+        return self.keys is None
+
+    @property
+    def nbytes(self):
+        if self.keys is None:
+            return 0
+        return self.keys.nbytes + self.values.nbytes
+
 
 class BatchRotatingKVCache(_BaseCache):
     step = 256
@@ -952,6 +1199,10 @@ class BatchRotatingKVCache(_BaseCache):
         self.offset += keys.shape[2]
         self._offset += keys.shape[2]
         self._idx = self.keys.shape[2]
+
+        # Make sure left_padding and offset are evaluated
+        self.keys = mx.depends(self.keys, (self.left_padding, self.offset))
+
         return self.keys, self.values
 
     def _update_in_place(self, keys, values):
@@ -1002,6 +1253,9 @@ class BatchRotatingKVCache(_BaseCache):
         self.offset += S
         self._idx += S
 
+        # Make sure left_padding and offset are evaluated
+        self.keys = mx.depends(self.keys, (self.left_padding, self.offset))
+
         # If the buffer is not full, slice off the end
         if self._offset < self.max_size:
             return (
@@ -1015,9 +1269,6 @@ class BatchRotatingKVCache(_BaseCache):
             return self._update_in_place(keys, values)
         return self._update_concat(keys, values)
 
-    def __len__(self):
-        return min(self._offset, self.max_size)
-
     def prepare(self, *, left_padding=None, lengths=None, right_padding=None):
         if left_padding is not None:
             if self.keys is not None:
@@ -1109,8 +1360,9 @@ class BatchRotatingKVCache(_BaseCache):
         """
         In-place filter to keep just the given indices in the cache.
         """
-        self.keys = self.keys[batch_indices]
-        self.values = self.values[batch_indices]
+        if self.keys is not None:
+            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]
 
@@ -1118,17 +1370,33 @@ class BatchRotatingKVCache(_BaseCache):
         """
         In-place extend this cache with the other cache.
         """
+        if self.keys is None and other.keys is None:
+            self.left_padding = mx.concatenate([self.left_padding, other.left_padding])
+            self.offset = mx.concatenate([self.offset, other.offset])
+            return
+
         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])
+        L1 = L2 = 0
+        if self.keys is not None:
+            B, H, L1, D = self.keys.shape
+            M = self.values.shape[3]
+        if other.keys is not None:
+            B, H, L2, D = other.keys.shape
+            M = other.values.shape[3]
+        max_size = max(L1, L2)
 
         def pad(c):
             left = max_idx - c._idx
-            right = max_size - c.keys.shape[2] - left
             k, v = c.keys, c.values
+            if k is None:
+                Bc = c.offset.shape[0]
+                k = mx.array([]).reshape(Bc, H, 0, D)
+                v = mx.array([]).reshape(Bc, H, 0, M)
+            right = max_size - k.shape[2] - left
             if right < 0:
                 k = k[..., :right, :]
                 v = v[..., :right, :]
@@ -1147,9 +1415,10 @@ class BatchRotatingKVCache(_BaseCache):
         self._offset = max(self._offset, other._offset)
 
     def extract(self, idx):
+        mx.eval(self.left_padding, self.offset)
         cache = RotatingKVCache(self.max_size)
-        padding = self.left_padding[idx].item()
-        offset = self.offset[idx].item()
+        padding = max(0, self.left_padding.tolist()[idx])
+        offset = self.offset.tolist()[idx]
         cache.keys = self.keys[idx : idx + 1]
         cache.values = self.values[idx : idx + 1]
         cache._idx = self._idx
@@ -1157,12 +1426,10 @@ class BatchRotatingKVCache(_BaseCache):
             cache.keys = mx.roll(cache.keys, -self._idx, axis=2)
             cache.values = mx.roll(cache.values, -self._idx, axis=2)
             cache._idx = self.max_size
-        if padding > 0:
-            cache.keys = mx.contiguous(cache.keys[:, :, padding : cache._idx])
-            cache.values = mx.contiguous(cache.values[:, :, padding : cache._idx])
+        cache.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
@@ -1173,8 +1440,13 @@ class BatchRotatingKVCache(_BaseCache):
             )
 
         offsets = [c.offset for c in caches]
-        lengths = [len(c) for c in caches]
+        lengths = [c.size() for c in caches]
         max_length = max(lengths)
+
+        # No cache has content so make an empty one
+        if max_length == 0:
+            return cls(caches[0].max_size, [0] * len(caches))
+
         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)
@@ -1184,9 +1456,11 @@ class BatchRotatingKVCache(_BaseCache):
 
         keys = mx.zeros((B, H, max_length, Dk), dtype=dt)
         values = mx.zeros((B, H, max_length, Dv), dtype=dt)
-        for i, (p, c) in enumerate(zip(padding, caches)):
-            keys[i : i + 1, :, p : p + c.offset] = c._temporal_order(c.keys)
-            values[i : i + 1, :, p : p + c.offset] = c._temporal_order(c.values)
+        for i, (p, l, c) in enumerate(zip(padding, lengths, caches)):
+            if c.keys is None:
+                continue
+            keys[i : i + 1, :, p : p + l] = c._temporal_order(c.keys)[..., -l:, :]
+            values[i : i + 1, :, p : p + l] = c._temporal_order(c.values)[..., -l:, :]
 
         cache = cls(caches[0].max_size, padding)
         cache.keys = keys
@@ -1196,3 +1470,294 @@ class BatchRotatingKVCache(_BaseCache):
         cache._offset = keys.shape[2]
 
         return cache
+
+    def size(self):
+        return min(self._offset, self.max_size)
+
+    def empty(self):
+        return self.keys is None
+
+    @property
+    def nbytes(self):
+        if self.keys is None:
+            return 0
+        return self.keys.nbytes + self.values.nbytes
+
+
+class TokenBuffer:
+    """A simple token buffer that can be efficiently appended to in a similar
+    fashion to the KVCache.
+
+    Perhaps these could share some logic in the future.
+    """
+
+    step = 256
+
+    def __init__(self, tokens=[]):
+        self._buffer = mx.array(tokens, dtype=mx.int32)
+        self._size = len(tokens)
+
+    def update_and_fetch(self, tokens):
+        start = self._size
+        end = start + len(tokens)
+
+        new_size = ((end + self.step - 1) // self.step) * self.step
+        if new_size > self._buffer.size:
+            self._buffer = mx.concatenate(
+                [self._buffer, mx.zeros(new_size - self._buffer.size, dtype=mx.int32)]
+            )
+        self._buffer[start:end] = tokens
+        self._size = end
+
+        return self._buffer[:end]
+
+    @property
+    def state(self):
+        return self._buffer
+
+    @property
+    def tokens(self):
+        return self._buffer[: self._size]
+
+
+@dataclass
+class PromptTrieResult:
+    model: Any
+    exact: Optional[List[int]]  # Exact match found
+    shorter: Optional[List[int]]  # Longest prefix with a value
+    longer: Optional[List[int]]  # Shortest value that extends beyond tokens
+    common_prefix: int  # Length of common prefix with any path
+
+
+class PromptTrie:
+    def __init__(self):
+        self._trie = {}
+
+    def add(self, model: Any, tokens: List[int], value: Any):
+        if model not in self._trie:
+            self._trie[model] = {}
+
+        current = self._trie[model]
+        for tok in tokens:
+            if tok not in current:
+                current[tok] = {}
+            current = current[tok]
+        prev = current.get("__value__", None)
+        current["__value__"] = value
+        return prev
+
+    def get(self, model: Any, tokens: List[int]):
+        current = self._trie[model]
+        for tok in tokens:
+            current = current[tok]
+        return current["__value__"]
+
+    def pop(self, model: Any, tokens: List[int]):
+        path = [self._trie[model]]
+        for tok in tokens:
+            path.append(path[-1][tok])
+        value = path[-1].pop("__value__")
+        for i in range(len(tokens), 0, -1):
+            node = path[i]
+            parent = path[i - 1]
+            tok = tokens[i - 1]
+            if len(node) > 0:
+                break
+            del parent[tok]
+        return value
+
+    def pop_prefixes(self, model: Any, tokens: List[int]):
+        values = []
+        current = self._trie[model]
+        for i, tok in enumerate(tokens):
+            if "__value__" in current:
+                values.append((i, current.pop("__value__")))
+            current = current[tok]
+        return values
+
+    def search(self, model: Any, tokens: List[int]) -> PromptTrieResult:
+        if model not in self._trie:
+            return PromptTrieResult(model, None, None, None, 0)
+
+        current = self._trie[model]
+
+        if not tokens and "__value__" in current:
+            return PromptTrieResult(model, [], None, None, 0)
+
+        # Walk the tokens as far as we can
+        last_index = -1
+        index = 0
+        while index < len(tokens) and tokens[index] in current:
+            current = current[tokens[index]]
+            if "__value__" in current:
+                last_index = index
+            index += 1
+
+        # Got an exact match
+        if last_index == len(tokens) - 1 >= 0:
+            return PromptTrieResult(model, tokens, None, None, 0)
+
+        # Check if we found a prefix at any point
+        shorter = None
+        if last_index > 0:
+            shorter = tokens[: last_index + 1]
+
+        # Check for sequences that are longer
+        longer = None
+        common_prefix = index
+        if index > 0:
+            best = None
+            stack = [(current, [])]
+            while stack:
+                current, extra = stack.pop()
+                if "__value__" in current:
+                    if best is None or len(extra) < len(best):
+                        best = extra
+                elif best is None or len(extra) < len(best):
+                    for tok in current:
+                        stack.append((current[tok], extra + [tok]))
+            longer = tokens[:index] + best
+        return PromptTrieResult(model, None, shorter, longer, common_prefix)
+
+
+class LRUPromptCache:
+    @dataclass
+    class CacheEntry:
+        prompt_cache: List[Any]
+        nbytes: int
+        cache_type: str
+
+    class CacheOrder:
+        def __init__(self, ordering: List[str] = ["assistant", "user", "system"]):
+            self._ordering = ordering
+            self._lrus = {k: deque() for k in ordering}
+
+        def __len__(self):
+            return sum(len(lru) for lru in self._lrus.values())
+
+        def push(self, model: Any, tokens: List[Any], cache_type: str = "assistant"):
+            self._lrus[cache_type].append((model, tokens))
+
+        def remove(self, model: Any, tokens: List[Any]):
+            for cache_type in self._ordering:
+                try:
+                    self._lrus[cache_type].remove((model, tokens))
+                    break
+                except ValueError:
+                    pass
+
+        def pop(self):
+            i = 0
+            while i + 1 < len(self._ordering):
+                lru_a = self._lrus[self._ordering[i]]
+                lru_b = self._lrus[self._ordering[i + 1]]
+                if lru_a and len(lru_a) >= len(lru_b):
+                    return lru_a.popleft()
+                i += 1
+            return lru_b.popleft()
+
+    def __init__(self, max_size: int = 10, max_bytes: int = 1 << 63):
+        self.max_size = max_size
+        self.max_bytes = max_bytes
+        self._trie = PromptTrie()
+        self._lru = LRUPromptCache.CacheOrder()
+        self._n_bytes = 0
+        self._n_bytes_by_type = {k: 0 for k in self._lru._ordering}
+
+    def __len__(self):
+        return len(self._lru)
+
+    @property
+    def nbytes(self):
+        return self._n_bytes
+
+    def fetch_nearest_cache(self, model: Any, tokens: List[int]):
+        result = self._trie.search(model, tokens)
+        if result.exact is not None:
+            cache_entry = self._trie.get(result.model, result.exact)
+            return copy.deepcopy(cache_entry.prompt_cache), []
+
+        short_length = len(result.shorter) if result.shorter is not None else 0
+        if result.longer is not None and result.common_prefix > short_length:
+            cache_entry = self._trie.get(result.model, result.longer)
+            if can_trim_prompt_cache(cache_entry.prompt_cache):
+                cache = copy.deepcopy(cache_entry.prompt_cache)
+                prefix = min(len(tokens) - 1, result.common_prefix)
+                num_to_trim = len(result.longer) - prefix
+                trim_prompt_cache(cache, num_to_trim)
+                return cache, tokens[prefix:]
+
+        if short_length > 0:
+            cache_entry = self._trie.get(result.model, result.shorter)
+            return copy.deepcopy(cache_entry.prompt_cache), tokens[short_length:]
+
+        return None, tokens
+
+    def insert_cache(
+        self,
+        model: Any,
+        tokens: List[int],
+        prompt_cache: List[Any],
+        *,
+        cache_type: str = "assistant",
+    ):
+        # Make the cache entry
+        entry = LRUPromptCache.CacheEntry(
+            prompt_cache, sum(c.nbytes for c in prompt_cache), cache_type
+        )
+
+        # Insert into the trie and update the byte counter and lru position
+        self._n_bytes += entry.nbytes
+        self._n_bytes_by_type[cache_type] += entry.nbytes
+        prev = self._trie.add(model, tokens, entry)
+        if prev is not None:
+            self._n_bytes -= prev.nbytes
+            self._n_bytes_by_type[prev.cache_type] -= prev.nbytes
+            self._lru.remove(model, tokens)
+        self._lru.push(model, tokens, cache_type)
+
+        # If it is a trimmable cache remove all prefixes cause they just take
+        # space
+        if can_trim_prompt_cache(prompt_cache):
+            for prefix_len, entry in self._trie.pop_prefixes(model, tokens):
+                self._n_bytes -= entry.nbytes
+                self._n_bytes_by_type[entry.cache_type] -= entry.nbytes
+                self._lru.remove(model, tokens[:prefix_len])
+
+        # Ensure we match the constraints
+        if len(self._lru) > self.max_size:
+            model, tokens = self._lru.pop()
+            entry = self._trie.pop(model, tokens)
+            self._n_bytes -= entry.nbytes
+            self._n_bytes_by_type[entry.cache_type] -= entry.nbytes
+        while self._n_bytes > self.max_bytes:
+            model, tokens = self._lru.pop()
+            entry = self._trie.pop(model, tokens)
+            self._n_bytes -= entry.nbytes
+            self._n_bytes_by_type[entry.cache_type] -= entry.nbytes
+
+    def trim_to(
+        self, *, n_sequences: Optional[int] = None, n_bytes: Optional[int] = None
+    ):
+        n_sequences = max(0, n_sequences) if n_sequences is not None else 1 << 63
+        n_bytes = max(0, n_bytes) if n_bytes is not None else 1 << 63
+
+        while len(self._lru) > n_sequences:
+            model, tokens = self._lru.pop()
+            entry = self._trie.pop(model, tokens)
+            self._n_bytes -= entry.nbytes
+            self._n_bytes_by_type[entry.cache_type] -= entry.nbytes
+        while self._n_bytes > n_bytes:
+            model, tokens = self._lru.pop()
+            entry = self._trie.pop(model, tokens)
+            self._n_bytes -= entry.nbytes
+            self._n_bytes_by_type[entry.cache_type] -= entry.nbytes
+
+    def stats_by_type(self):
+        result = {}
+        for cache_type in self._lru._ordering:
+            result[cache_type] = {
+                "n_sequences": len(self._lru._lrus[cache_type]),
+                "n_bytes": self._n_bytes_by_type[cache_type],
+            }
+        return result

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,7 +6,9 @@ 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
@@ -259,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
 
 
@@ -315,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
 
 
@@ -395,7 +405,8 @@ class DeepseekV2Model(PipelineMixin, nn.Module):
                 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)
 
@@ -429,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.pipeline_layers

mlx_lm/models/deepseek_v3.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 .mla import MultiLinear
 from .pipeline import PipelineMixin
 from .rope_utils import initialize_rope
 from .switch_layers import SwitchGLU
@@ -83,11 +86,11 @@ class DeepseekV3Attention(nn.Module):
             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.embed_q = MultiLinear(
+            self.qk_nope_head_dim, self.kv_lora_rank, self.num_heads
+        )
+        self.unembed_out = MultiLinear(
+            self.kv_lora_rank, self.v_head_dim, self.num_heads
         )
 
         self.o_proj = nn.Linear(
@@ -130,29 +133,38 @@ class DeepseekV3Attention(nn.Module):
         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)
+        kv_latent = self.kv_a_layernorm(compressed_kv)
 
-        k_nope, values = mx.split(kv, [self.qk_nope_head_dim], axis=-1)
+        offset = cache.offset if cache is not None else 0
+        q_pe = self.rope(q_pe, offset)
+        k_pe = self.rope(k_pe, offset)
+
+        kv_latent = mx.expand_dims(kv_latent, 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)
+            kv_latent, k_pe = cache.update_and_fetch(kv_latent, k_pe)
 
-        queries = mx.concatenate([q_nope, q_pe], axis=-1)
+        pe_scores = (q_pe * self.scale) @ k_pe.swapaxes(-1, -2)
+        if mask is not None:
+            pe_scores = mx.where(
+                mask,
+                pe_scores,
+                mx.array(mx.finfo(pe_scores.dtype).min, pe_scores.dtype),
+            )
+
+        if L == 1:
+            q_nope = self.embed_q(q_nope)
+            k = v = kv_latent
+        else:
+            k = self.embed_q(kv_latent, transpose=False)
+            v = self.unembed_out(kv_latent)
 
         output = scaled_dot_product_attention(
-            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+            q_nope, k, v, cache=cache, scale=self.scale, mask=pe_scores
         )
+        if L == 1:
+            output = self.unembed_out(output)
+
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
         return self.o_proj(output)
 
@@ -173,7 +185,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
 
 
@@ -256,13 +268,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
 
 
@@ -319,7 +339,7 @@ class DeepseekV3Model(PipelineMixin, nn.Module):
 
         if cache is None:
             cache = [None] * len(self.pipeline_layers)
-        mask = create_attention_mask(h, cache[0])
+        mask = create_attention_mask(h, cache[0], return_array=True)
 
         # Receive from the previous process in the pipeline
         if pipeline_rank < pipeline_size - 1:
@@ -335,7 +355,8 @@ class DeepseekV3Model(PipelineMixin, nn.Module):
                 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)
 
@@ -358,7 +379,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
@@ -411,6 +433,42 @@ class Model(nn.Module):
                             for e in range(self.args.n_routed_experts)
                         ]
                         weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
+            prefix = f"model.layers.{l}.self_attn"
+            if f"{prefix}.kv_b_proj.weight" in weights:
+                layer = self.model.layers[l].self_attn.embed_q
+                quantized = f"{prefix}.kv_b_proj.scales" in weights
+                v = weights.pop(f"{prefix}.kv_b_proj.weight")
+                head_dim = self.args.qk_nope_head_dim + self.args.v_head_dim
+
+                if quantized:
+                    dims = self.args.kv_lora_rank
+                    scales = weights.pop(f"{prefix}.kv_b_proj.scales")
+                    biases = weights.pop(f"{prefix}.kv_b_proj.biases")
+                    # Try to infer bits and group size
+                    bits = (v.shape[-1] * 32) // dims
+                    group_size = dims // scales.shape[-1]
+                    v = mx.dequantize(
+                        v, scales, biases, bits=bits, group_size=group_size
+                    )
+                num_heads = self.args.num_attention_heads
+                v = v.reshape(num_heads, head_dim, -1)
+                wk = mx.contiguous(
+                    v[:, : self.args.qk_nope_head_dim, :].swapaxes(-1, -2)
+                )
+                wv = mx.contiguous(v[:, self.args.qk_nope_head_dim :, :])
+                if quantized:
+                    wk, wk_scales, wk_biases = mx.quantize(
+                        wk, bits=bits, group_size=group_size
+                    )
+                    wv, wv_scales, wv_biases = mx.quantize(
+                        wv, bits=bits, group_size=group_size
+                    )
+                    weights[f"{prefix}.embed_q.scales"] = wk_scales
+                    weights[f"{prefix}.unembed_out.scales"] = wv_scales
+                    weights[f"{prefix}.embed_q.biases"] = wk_biases
+                    weights[f"{prefix}.unembed_out.biases"] = wv_biases
+                weights[f"{prefix}.embed_q.weight"] = wk
+                weights[f"{prefix}.unembed_out.weight"] = wv
 
         # Remove multi-token prediction layer and any unused precomputed rotary freqs
         return {
@@ -419,6 +477,70 @@ 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()
+        rank = group.rank()
+        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.num_heads //= N
+            num_heads = layer.self_attn.num_heads
+            sh = rank * num_heads
+            eh = sh + num_heads
+
+            def shard_heads(w):
+                return w[sh:eh]
+
+            layer.self_attn.embed_q.apply(shard_heads)
+            layer.self_attn.unembed_out.apply(shard_heads)
+
+            layer.self_attn.o_proj = shard_linear(
+                layer.self_attn.o_proj, "sharded-to-all", group=group
+            )
+
+            # 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.pipeline_layers

mlx_lm/models/deepseek_v32.py

@@ -6,9 +6,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 .cache import CacheList, KVCache
+from .mla import MultiLinear
 from .rope_utils import initialize_rope
 from .switch_layers import SwitchGLU
 
@@ -68,7 +71,7 @@ class Indexer(nn.Module):
         self.rope = initialize_rope(
             dims=args.qk_rope_head_dim,
             base=args.rope_theta,
-            traditional=False,
+            traditional=True,
             max_position_embeddings=args.max_position_embeddings,
             scaling_config=args.rope_scaling,
         )
@@ -84,19 +87,15 @@ class Indexer(nn.Module):
         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)
+
+        offset = cache.offset if cache is not None else 0
+
+        q = self.rope(q, offset=offset)
+        k = self.rope(k, offset=offset)
+
         if cache is not None:
             k, _ = cache.update_and_fetch(k, mx.zeros([b, 1, s, 0]))
         if k.shape[2] <= self.index_topk:
@@ -106,7 +105,7 @@ class Indexer(nn.Module):
         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)
+        scores = scores.sum(axis=1, keepdims=True)
         if mask is not None:
             scores = mx.where(mask, scores, -float("inf"))
         return mx.argpartition(scores, kth=-self.index_topk, axis=-1)[
@@ -145,11 +144,11 @@ class DeepseekV32Attention(nn.Module):
             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.embed_q = MultiLinear(
+            self.qk_nope_head_dim, self.kv_lora_rank, self.num_heads
+        )
+        self.unembed_out = MultiLinear(
+            self.kv_lora_rank, self.v_head_dim, self.num_heads
         )
 
         self.o_proj = nn.Linear(
@@ -191,40 +190,71 @@ class DeepseekV32Attention(nn.Module):
         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)
+        kv_latent = self.kv_a_layernorm(compressed_kv)
 
-        k_nope, values = mx.split(kv, [self.qk_nope_head_dim], axis=-1)
+        offset = cache[0].offset if cache is not None else 0
+        q_pe = self.rope(q_pe, offset)
+        k_pe = self.rope(k_pe, offset)
+
+        kv_latent = mx.expand_dims(kv_latent, 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
-            )
+            kv_latent, k_pe = cache[0].update_and_fetch(kv_latent, k_pe)
         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
+            if L == 1:
+                idx = topk_indices[:, :, 0, :, None]
+                kv_latent = mx.take_along_axis(
+                    kv_latent,
+                    mx.broadcast_to(idx, idx.shape[:-1] + (kv_latent.shape[-1],)),
+                    axis=2,
+                )
+                k_pe = mx.take_along_axis(
+                    k_pe,
+                    mx.broadcast_to(idx, idx.shape[:-1] + (k_pe.shape[-1],)),
+                    axis=2,
+                )
+                if mask is not None:
+                    mask = mx.take_along_axis(mask, topk_indices, axis=-1)
+            else:
+                shape = list(topk_indices.shape)
+                shape[-1] = kv_latent.shape[2]
+                sparse_mask = mx.zeros(shape, dtype=mx.bool_)
+                sparse_mask = mx.put_along_axis(
+                    sparse_mask, topk_indices, mx.array(True), axis=-1
+                )
+                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))
+
+        pe_scores = (q_pe * self.scale) @ k_pe.swapaxes(-1, -2)
+        if mask is not None:
+            pe_scores = mx.where(
+                mask,
+                pe_scores,
+                mx.array(mx.finfo(pe_scores.dtype).min, pe_scores.dtype),
             )
-            sparse_mask = sparse_mask[:, None, :, :]
-            if mask is not None:
-                sparse_mask = sparse_mask & mask
-            mask = sparse_mask
+
+        if L == 1:
+            q_nope = self.embed_q(q_nope)
+            k = v = kv_latent
+        else:
+            k = self.embed_q(kv_latent, transpose=False)
+            v = self.unembed_out(kv_latent)
+
         output = scaled_dot_product_attention(
-            queries, keys, values, cache=cache[0], scale=self.scale, mask=mask
+            q_nope, k, v, cache=cache, scale=self.scale, mask=pe_scores
         )
+        if L == 1:
+            output = self.unembed_out(output)
+
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
         return self.o_proj(output)
 
@@ -245,7 +275,7 @@ class DeepseekV32MLP(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
 
 
@@ -328,13 +358,21 @@ class DeepseekV32MoE(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
 
 
@@ -428,10 +466,11 @@ class DeepseekV32Model(nn.Module):
         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)
+                cache[-1][0].keys = mx.depends(cache[-1][0].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)
 
@@ -453,8 +492,19 @@ class Model(nn.Module):
         return self.lm_head(out)
 
     def sanitize(self, weights):
+        # Remove multi-token prediction layers
+        mpt_layer = self.args.num_hidden_layers
+        new_weights = {}
+        for k, v in weights.items():
+            parts = k.split(".")
+            if len(parts) >= 3 and parts[1] == "layers" and int(parts[2]) >= mpt_layer:
+                continue
+            new_weights[k] = v
+        weights = new_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
@@ -492,13 +542,102 @@ class Model(nn.Module):
                             for e in range(self.args.n_routed_experts)
                         ]
                         weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
+            prefix = f"model.layers.{l}.self_attn"
+            if f"{prefix}.kv_b_proj.weight" in weights:
+                layer = self.model.layers[l].self_attn.embed_q
+                quantized = f"{prefix}.kv_b_proj.scales" in weights
+                v = weights.pop(f"{prefix}.kv_b_proj.weight")
+                head_dim = self.args.qk_nope_head_dim + self.args.v_head_dim
 
-        # 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
-        }
+                if quantized:
+                    dims = self.args.kv_lora_rank
+                    scales = weights.pop(f"{prefix}.kv_b_proj.scales")
+                    biases = weights.pop(f"{prefix}.kv_b_proj.biases")
+                    # Try to infer bits and group size
+                    bits = (v.shape[-1] * 32) // dims
+                    group_size = dims // scales.shape[-1]
+                    v = mx.dequantize(
+                        v, scales, biases, bits=bits, group_size=group_size
+                    )
+                num_heads = self.args.num_attention_heads
+                v = v.reshape(num_heads, head_dim, -1)
+                wk = mx.contiguous(
+                    v[:, : self.args.qk_nope_head_dim, :].swapaxes(-1, -2)
+                )
+                wv = mx.contiguous(v[:, self.args.qk_nope_head_dim :, :])
+                if quantized:
+                    wk, wk_scales, wk_biases = mx.quantize(
+                        wk, bits=bits, group_size=group_size
+                    )
+                    wv, wv_scales, wv_biases = mx.quantize(
+                        wv, bits=bits, group_size=group_size
+                    )
+                    weights[f"{prefix}.embed_q.scales"] = wk_scales
+                    weights[f"{prefix}.unembed_out.scales"] = wv_scales
+                    weights[f"{prefix}.embed_q.biases"] = wk_biases
+                    weights[f"{prefix}.unembed_out.biases"] = wv_biases
+                weights[f"{prefix}.embed_q.weight"] = wk
+                weights[f"{prefix}.unembed_out.weight"] = wv
+
+        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:
+            layer.self_attn.q_b_proj = shard_linear(
+                layer.self_attn.q_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
+            num_heads = layer.self_attn.num_heads
+            sh = rank * num_heads
+            eh = sh + num_heads
+
+            def shard_heads(w):
+                return w[sh:eh]
+
+            layer.self_attn.embed_q.apply(shard_heads)
+            layer.self_attn.unembed_out.apply(shard_heads)
+
+            # 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):

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

@@ -6,13 +6,14 @@ 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 .cache import ArraysCache, CacheList, KVCache
 from .rope_utils import initialize_rope
 from .ssm import ssm_update
 
@@ -69,6 +70,7 @@ class ModelArgs(BaseModelArgs):
     )
     ssm_out_multiplier: float = 0.23570226039551587
     vocab_size: int = 32784
+    tie_word_embeddings: bool = True
 
 
 class FalconH1RMSNormGated(nn.Module):
@@ -81,14 +83,14 @@ class FalconH1RMSNormGated(nn.Module):
 
     def __call__(self, hidden_states, gate=None):
         if not self.norm_before_gate and gate is not None:
-            hidden_states = hidden_states * nn.silu(gate)
+            hidden_states = 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 = hidden_states * nn.silu(gate)
+            hidden_states = swiglu(gate, hidden_states)
         return hidden_states
 
 
@@ -231,21 +233,36 @@ class FalconH1Mixer(nn.Module):
             self.intermediate_size, self.hidden_size, bias=args.projectors_bias
         )
 
-    def _apply_conv(
-        self, conv_input: mx.array, cache: Optional[MambaCache] = None
+    def _conv(
+        self,
+        conv_input: mx.array,
+        cache: Optional[ArraysCache],
+        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)
@@ -256,17 +273,20 @@ class FalconH1Mixer(nn.Module):
         B: mx.array,
         C: mx.array,
         dt: mx.array,
-        state: Optional[mx.array] = None,
-        mask: Optional[mx.array] = None,
+        cache: Optional[ArraysCache],
+        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,
@@ -278,9 +298,11 @@ class FalconH1Mixer(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, input_states, cache=None, mask: Optional[mx.array] = None):
         projected_states = self.in_proj(input_states)
@@ -291,11 +313,9 @@ class FalconH1Mixer(nn.Module):
             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(
+        hidden_states, B, C = mx.split(
             conv_output,
             [
                 self.intermediate_size,
@@ -303,15 +323,15 @@ class FalconH1Mixer(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, B, C, dt, cache, mask=mask)
         if cache:
-            cache[1] = state
+            cache.advance(y.shape[1])
 
         if self.mamba_rms_norm:
             y = self.norm(y, gate)
         else:
-            y = y * nn.silu(gate)
+            y = swiglu(gate, y)
 
         return self.out_proj(y)
 
@@ -329,7 +349,7 @@ class FalconH1MLP(nn.Module):
         self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=args.mlp_bias)
 
     def __call__(self, x):
-        y = self.up_proj(x) * nn.silu(self.gate_proj(x))
+        y = swiglu(self.gate_proj(x), self.up_proj(x))
         y = self.down_proj(y)
         return y
 
@@ -425,11 +445,16 @@ class Model(nn.Module):
         self.args = args
         self.model_type = args.model_type
         self.model = FalconH1Model(args=args)
-        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+        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)
-        return self.lm_head(hidden_states)
+        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
@@ -470,7 +495,7 @@ class Model(nn.Module):
 
     def make_cache(self):
         return [
-            CacheList(MambaCache(), KVCache())
+            CacheList(ArraysCache(size=2), KVCache())
             for _ in range(self.args.num_hidden_layers)
         ]
 

mlx_lm/models/gated_delta.py

@@ -7,9 +7,7 @@ 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))
 
 
 def _make_gated_delta_kernel(has_mask=False, vectorized=False):
@@ -83,6 +81,8 @@ def _make_gated_delta_kernel(has_mask=False, vectorized=False):
             if (thread_index_in_simdgroup == 0) {{
               y[dv_idx] = static_cast<InT>(out);
             }}
+          }} else {{
+            y[dv_idx] = static_cast<InT>(0);
           }}
           // Increment data pointers to next time step
           q_ += Hk * Dk;
@@ -94,7 +94,7 @@ def _make_gated_delta_kernel(has_mask=False, vectorized=False):
         }}
         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]);
+          o_state[s_idx] = static_cast<StT>(state[i]);
         }}
     """
     inputs = ["q", "k", "v", "g", "beta", "state_in", "T"]
@@ -161,13 +161,11 @@ def _gated_delta_step_ops(
     state = state + k[..., None, :] * delta[..., None]
     # Output projection along key dim with q
     y = (state * q[..., None, :]).sum(axis=-1)  # [B, H, Dv]
+
     if mask is not None:
-        if mask.ndim == 2:
-            mask = mx.expand_dims(mask, axes=(2, 3))
-        elif mask.ndim == 3:
-            mask = mx.expand_dims(mask, axis=-1)
+        mask = mx.expand_dims(mask, axis=(1, 2, 3))
         state = mx.where(mask, state, old_state)
-    return y, state
+    return y.astype(q.dtype), state
 
 
 def gated_delta_kernel(
@@ -182,6 +180,7 @@ def gated_delta_kernel(
     B, T, Hk, Dk = k.shape
     Hv, Dv = v.shape[2:]
     input_type = q.dtype
+    state_type = state.dtype
     if g.ndim == 4:
         kernel = _gated_delta_kernel_vec
         inputs = [q, k, v, g, beta, state, T]
@@ -199,6 +198,7 @@ def gated_delta_kernel(
         inputs=inputs,
         template=[
             ("InT", input_type),
+            ("StT", state_type),
             ("Dk", Dk),
             ("Dv", Dv),
             ("Hk", Hk),
@@ -207,7 +207,7 @@ def gated_delta_kernel(
         grid=(32, Dv, B * Hv),
         threadgroup=(32, 4, 1),
         output_shapes=[(B, T, Hv, Dv), state.shape],
-        output_dtypes=[input_type, input_type],
+        output_dtypes=[input_type, state_type],
     )
 
 
@@ -237,7 +237,7 @@ def gated_delta_ops(
     B, T, Hk, Dk = q.shape
     Hv, Dv = v.shape[-2:]
     if state is None:
-        state = mx.zeros((B, Hv, Dv, Dk), dtype=q.dtype)
+        state = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32)
 
     if (repeat_factor := Hv // Hk) > 1:
         q = mx.repeat(q, repeat_factor, -2)
@@ -271,13 +271,12 @@ def gated_delta_update(
     mask: Optional[mx.array] = None,
     use_kernel: bool = True,
 ) -> Tuple[mx.array, mx.array]:
-
     beta = mx.sigmoid(b)
     g = compute_g(A_log, a, dt_bias)
     if state is None:
         B, _, Hk, Dk = q.shape
         Hv, Dv = v.shape[-2:]
-        state = mx.zeros((B, Hv, Dv, Dk), dtype=q.dtype)
+        state = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32)
 
     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)

mlx_lm/models/gemma3_text.py

@@ -54,13 +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 = initialize_rope(
-            dims=head_dim,
-            base=(args.rope_local_base_freq if self.is_sliding else args.rope_theta),
-            traditional=False,
-            max_position_embeddings=args.max_position_embeddings,
-            scaling_config=args.rope_scaling,
-        )
+        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,

mlx_lm/models/gemma4.py

@@ -0,0 +1,92 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten, tree_unflatten
+
+from . import gemma4_text
+from .base import BaseModelArgs
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str = "gemma4"
+    text_config: dict = None
+    vocab_size: int = 262144
+
+    def __post_init__(self):
+        if self.text_config is None:
+            self.text_config = {}
+        self.text_config["vocab_size"] = self.vocab_size
+        self.text_config["num_attention_heads"] = self.text_config.get(
+            "num_attention_heads", 8
+        )
+        self.text_config["num_key_value_heads"] = self.text_config.get(
+            "num_key_value_heads", 1
+        )
+
+
+class Model(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.args = args
+        self.model_type = args.model_type
+        self.language_model = gemma4_text.Model(
+            gemma4_text.ModelArgs.from_dict(args.text_config)
+        )
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache=None,
+        input_embeddings: Optional[mx.array] = None,
+        per_layer_inputs: Optional[mx.array] = None,
+    ):
+        return self.language_model(
+            inputs,
+            cache=cache,
+            input_embeddings=input_embeddings,
+            per_layer_inputs=per_layer_inputs,
+        )
+
+    def sanitize(self, weights):
+        new_weights = {}
+        for k, v in weights.items():
+            starts_w_model = k.startswith("model.")
+
+            k = k.removeprefix("model.")
+            if k.startswith(
+                (
+                    "vision_tower",
+                    "multi_modal_projector",
+                    "audio_tower",
+                    "embed_audio",
+                    "embed_vision",
+                )
+            ):
+                continue
+
+            if not starts_w_model:
+                new_weights[k] = v
+                continue
+
+            if k.startswith("language_model"):
+                k = k.replace("language_model.", "language_model.model.")
+
+            new_weights[k] = v
+
+        return self.language_model.sanitize(new_weights)
+
+    @property
+    def layers(self):
+        return self.language_model.layers
+
+    @property
+    def quant_predicate(self):
+        return self.language_model.quant_predicate
+
+    def make_cache(self):
+        return self.language_model.make_cache()

mlx_lm/models/gemma4_text.py

@@ -0,0 +1,688 @@
+# Copyright © 2025 Apple Inc.
+
+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 .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
+from .cache import KVCache, RotatingKVCache, _BaseCache
+from .rope_utils import initialize_rope
+from .switch_layers import SwitchGLU
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str = "gemma4_text"
+    hidden_size: int = 1536
+    num_hidden_layers: int = 35
+    intermediate_size: int = 6144
+    num_attention_heads: int = 8
+    head_dim: int = 256
+    global_head_dim: int = 512
+    global_partial_rotary_factor: float = 0.25
+    rms_norm_eps: float = 1e-6
+    vocab_size: int = 262144
+    vocab_size_per_layer_input: int = 262144
+    num_key_value_heads: int = 1
+    num_global_key_value_heads: Optional[int] = None
+    num_kv_shared_layers: int = 20
+    pad_token_id: int = 0
+    hidden_size_per_layer_input: int = 256
+    rope_traditional: bool = False
+    partial_rotary_factor: float = 1.0
+    rope_parameters: Optional[Dict] = None
+    sliding_window: int = 512
+    sliding_window_pattern: int = 5
+    max_position_embeddings: int = 131072
+    attention_k_eq_v: bool = False
+    final_logit_softcapping: float = 30.0
+    use_double_wide_mlp: bool = True
+    enable_moe_block: bool = False
+    num_experts: Optional[int] = None
+    top_k_experts: Optional[int] = None
+    moe_intermediate_size: Optional[int] = None
+    layer_types: Optional[List[str]] = None
+    tie_word_embeddings: bool = True
+
+    def __post_init__(self):
+        if self.rope_parameters is None:
+            self.rope_parameters = {
+                "full_attention": {
+                    "partial_rotary_factor": 0.25,
+                    "rope_theta": 1000000.0,
+                    "rope_type": "proportional",
+                },
+                "sliding_attention": {
+                    "partial_rotary_factor": 1.0,
+                    "rope_theta": 10000.0,
+                    "rope_type": "default",
+                },
+            }
+        if self.layer_types is None:
+            pattern = ["sliding_attention"] * (self.sliding_window_pattern - 1) + [
+                "full_attention"
+            ]
+            self.layer_types = (pattern * (self.num_hidden_layers // len(pattern) + 1))[
+                : self.num_hidden_layers
+            ]
+
+
+class RMSNormNoScale(nn.Module):
+    """RMSNorm without learnable scale."""
+
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return mx.fast.rms_norm(x, None, self.eps)
+
+
+@partial(mx.compile, shapeless=True)
+def logit_softcap(softcap, x):
+    return mx.tanh(x / softcap) * softcap
+
+
+@partial(mx.compile, shapeless=True)
+def _complete_square(x2, y2, xy):
+    return x2 + mx.expand_dims(y2, -1) - 2 * xy
+
+
+@partial(mx.compile, shapeless=True)
+def geglu(gate, x):
+    return nn.gelu_approx(gate) * x
+
+
+class MLP(nn.Module):
+    def __init__(self, config: ModelArgs, layer_idx: int = 0):
+        super().__init__()
+        first_kv_shared_layer_idx = (
+            config.num_hidden_layers - config.num_kv_shared_layers
+        )
+        is_kv_shared_layer = layer_idx >= first_kv_shared_layer_idx > 0
+        use_double_wide = config.use_double_wide_mlp and is_kv_shared_layer
+        intermediate_size = config.intermediate_size * (2 if use_double_wide else 1)
+
+        self.gate_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, config.hidden_size, bias=False)
+        self.up_proj = nn.Linear(config.hidden_size, intermediate_size, bias=False)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(geglu(self.gate_proj(x), self.up_proj(x)))
+
+
+class Router(nn.Module):
+    """Expert router: norm -> scale -> project -> top-k -> renormalize."""
+
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.eps = config.rms_norm_eps
+        self.proj = nn.Linear(config.hidden_size, config.num_experts, bias=False)
+        self.scale = mx.ones((config.hidden_size,))
+        self.per_expert_scale = mx.ones((config.num_experts,))
+        self._root_size = config.hidden_size**-0.5
+
+    def __call__(self, x: mx.array):
+        x = mx.fast.rms_norm(x, self.scale * self._root_size, self.eps)
+
+        expert_scores = self.proj(x)
+
+        top_k_indices = mx.argpartition(
+            expert_scores, kth=-self.config.top_k_experts, axis=-1
+        )
+        top_k_indices = top_k_indices[..., -self.config.top_k_experts :]
+
+        top_k_weights = mx.take_along_axis(expert_scores, top_k_indices, axis=-1)
+        top_k_weights = mx.softmax(top_k_weights, axis=-1)
+        top_k_weights = top_k_weights * self.per_expert_scale[top_k_indices]
+
+        return top_k_indices, top_k_weights
+
+
+class GeGLU(nn.Module):
+    """GELU-gated linear unit activation for SwitchGLU."""
+
+    def __call__(self, x, gate):
+        return geglu(gate, x)
+
+
+class Experts(nn.Module):
+    """Sparse MoE using SwitchGLU with gather_mm."""
+
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+
+        self.switch_glu = SwitchGLU(
+            input_dims=config.hidden_size,
+            hidden_dims=config.moe_intermediate_size,
+            num_experts=config.num_experts,
+            activation=GeGLU(),
+            bias=False,
+        )
+
+    def __call__(
+        self, x: mx.array, top_k_indices: mx.array, top_k_weights: mx.array
+    ) -> mx.array:
+        w = mx.expand_dims(top_k_weights, -1)
+        y = self.switch_glu(x, top_k_indices)
+
+        return (w * y).sum(-2)
+
+
+class Attention(nn.Module):
+    def __init__(self, config: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.layer_type = config.layer_types[layer_idx]
+        self.is_sliding = self.layer_type == "sliding_attention"
+        self.has_kv = layer_idx < config.num_hidden_layers - config.num_kv_shared_layers
+
+        self.head_dim = (
+            config.global_head_dim
+            if self.layer_type == "full_attention"
+            and hasattr(config, "global_head_dim")
+            and config.global_head_dim
+            else config.head_dim
+        )
+
+        dim = config.hidden_size
+        self.n_heads = config.num_attention_heads
+
+        # K-eq-V for full attention layers (26B/31B models)
+        self.use_k_eq_v = config.attention_k_eq_v and not self.is_sliding
+        if self.use_k_eq_v and config.num_global_key_value_heads is not None:
+            self.n_kv_heads = config.num_global_key_value_heads
+        else:
+            self.n_kv_heads = config.num_key_value_heads
+
+        self.scale = 1.0
+
+        self.q_proj = nn.Linear(dim, self.n_heads * self.head_dim, bias=False)
+        if self.has_kv:
+            self.k_proj = nn.Linear(dim, self.n_kv_heads * self.head_dim, bias=False)
+            if not self.use_k_eq_v:
+                self.v_proj = nn.Linear(
+                    dim, self.n_kv_heads * self.head_dim, bias=False
+                )
+        self.o_proj = nn.Linear(self.n_heads * self.head_dim, dim, bias=False)
+
+        self.q_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        if self.has_kv:
+            self.k_norm = nn.RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+            self.v_norm = RMSNormNoScale(self.head_dim, eps=config.rms_norm_eps)
+
+        # RoPE (with partial rotation support)
+        layer_key = "sliding_attention" if self.is_sliding else "full_attention"
+        rope_params = config.rope_parameters.get(layer_key, {})
+        rope_theta = rope_params.get("rope_theta", 10000.0)
+        self.rope = initialize_rope(
+            dims=self.head_dim,
+            traditional=config.rope_traditional,
+            base=rope_theta,
+            scaling_config=rope_params,
+            max_position_embeddings=config.max_position_embeddings,
+        )
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+        shared_kv: Optional[tuple] = None,
+        offset: Optional[Any] = None,
+    ) -> mx.array:
+        B, L, _ = x.shape
+
+        queries = self.q_proj(x).reshape(B, L, self.n_heads, self.head_dim)
+        queries = self.q_norm(queries)
+
+        if shared_kv is not None:
+            keys, values = shared_kv
+        elif not self.has_kv:
+            raise ValueError(
+                f"Layer {self.layer_idx} is a KV-shared layer but received no shared_kv"
+            )
+        else:
+            keys = self.k_proj(x).reshape(B, L, self.n_kv_heads, self.head_dim)
+            values = keys
+            if not self.use_k_eq_v:
+                values = self.v_proj(x).reshape(B, L, self.n_kv_heads, self.head_dim)
+
+            offset = mx.array(cache.offset) if cache is not None else 0
+
+            keys = self.k_norm(keys)
+            keys = keys.transpose(0, 2, 1, 3)
+            keys = self.rope(keys, offset=offset)
+
+            values = self.v_norm(values)
+            values = values.transpose(0, 2, 1, 3)
+
+        queries = queries.transpose(0, 2, 1, 3)
+        queries = self.rope(queries, offset=offset)
+
+        if cache is not None:
+            keys, values = cache.update_and_fetch(keys, values)
+
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
+        )
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+
+        return self.o_proj(output), (keys, values), offset
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, config: ModelArgs, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.layer_type = config.layer_types[layer_idx]
+        self.self_attn = Attention(config, layer_idx)
+        self.mlp = MLP(config, layer_idx)
+        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
+        )
+        self.pre_feedforward_layernorm = nn.RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+        self.post_feedforward_layernorm = nn.RMSNorm(
+            config.hidden_size, eps=config.rms_norm_eps
+        )
+
+        # MoE (26B model)
+        self.enable_moe = config.enable_moe_block
+        if self.enable_moe:
+            self.router = Router(config)
+            self.experts = Experts(config)
+            self.post_feedforward_layernorm_1 = nn.RMSNorm(
+                config.hidden_size, eps=config.rms_norm_eps
+            )
+            self.post_feedforward_layernorm_2 = nn.RMSNorm(
+                config.hidden_size, eps=config.rms_norm_eps
+            )
+            self.pre_feedforward_layernorm_2 = nn.RMSNorm(
+                config.hidden_size, eps=config.rms_norm_eps
+            )
+
+        # Per-layer input gating (2B/4B models)
+        self.hidden_size_per_layer_input = config.hidden_size_per_layer_input
+        if self.hidden_size_per_layer_input:
+            self.per_layer_input_gate = nn.Linear(
+                config.hidden_size, self.hidden_size_per_layer_input, bias=False
+            )
+            self.per_layer_projection = nn.Linear(
+                self.hidden_size_per_layer_input, config.hidden_size, bias=False
+            )
+            self.post_per_layer_input_norm = nn.RMSNorm(
+                config.hidden_size, eps=config.rms_norm_eps
+            )
+        else:
+            self.per_layer_input_gate = None
+            self.per_layer_projection = None
+            self.post_per_layer_input_norm = None
+
+        # Layer scalar
+        self.layer_scalar = mx.ones((1,))
+
+    def __call__(
+        self,
+        x: mx.array,
+        mask: Optional[mx.array] = None,
+        cache: Optional[Any] = None,
+        per_layer_input: Optional[mx.array] = None,
+        shared_kv: Optional[tuple] = None,
+        offset: Optional[Any] = None,
+    ) -> mx.array:
+        residual = x
+
+        h = self.input_layernorm(x)
+        h, shared_kv, offset = self.self_attn(
+            h, mask, cache, shared_kv=shared_kv, offset=offset
+        )
+        h = self.post_attention_layernorm(h)
+        h = residual + h
+
+        residual = h
+
+        if self.enable_moe:
+            h1 = self.pre_feedforward_layernorm(h)
+            h1 = self.mlp(h1)
+            h1 = self.post_feedforward_layernorm_1(h1)
+
+            top_k_indices, top_k_weights = self.router(h)
+            h2 = self.pre_feedforward_layernorm_2(h)
+            h2 = self.experts(h2, top_k_indices, top_k_weights)
+            h2 = self.post_feedforward_layernorm_2(h2)
+
+            h = h1 + h2
+        else:
+            h = self.pre_feedforward_layernorm(h)
+            h = self.mlp(h)
+
+        h = self.post_feedforward_layernorm(h)
+        h = residual + h
+
+        # Per-layer input gating
+        if (
+            self.per_layer_input_gate is not None
+            and self.per_layer_projection is not None
+            and self.post_per_layer_input_norm is not None
+            and per_layer_input is not None
+        ):
+            residual = h
+            gate = self.per_layer_input_gate(h)
+            gate = nn.gelu_approx(gate)
+            gate = mx.multiply(gate, per_layer_input)
+            gate = self.per_layer_projection(gate)
+            gate = self.post_per_layer_input_norm(gate)
+            h = residual + gate
+
+        if self.layer_scalar is not None:
+            h = h * self.layer_scalar
+
+        return h, shared_kv, offset
+
+
+class Gemma4TextModel(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.config = config
+        self.vocab_size = config.vocab_size
+        self.window_size = config.sliding_window
+        self.sliding_window_pattern = config.sliding_window_pattern
+        self.num_hidden_layers = config.num_hidden_layers
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.embed_scale = config.hidden_size**0.5
+        self.layers = [
+            DecoderLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)
+        ]
+        self.norm = nn.RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+        # Per-layer input embeddings (2B/4B models)
+        self.hidden_size_per_layer_input = config.hidden_size_per_layer_input
+        if self.hidden_size_per_layer_input:
+            self.embed_tokens_per_layer = nn.Embedding(
+                config.vocab_size_per_layer_input,
+                config.num_hidden_layers * config.hidden_size_per_layer_input,
+            )
+            self.embed_tokens_per_layer_scale = config.hidden_size_per_layer_input**0.5
+            self.per_layer_input_scale = 2.0**-0.5
+            self.per_layer_projection_scale = config.hidden_size**-0.5
+            self.per_layer_model_projection = nn.Linear(
+                config.hidden_size,
+                config.num_hidden_layers * config.hidden_size_per_layer_input,
+                bias=False,
+            )
+            self.per_layer_projection_norm = nn.RMSNorm(
+                config.hidden_size_per_layer_input, eps=config.rms_norm_eps
+            )
+        else:
+            self.embed_tokens_per_layer = None
+            self.per_layer_input_scale = None
+            self.per_layer_projection_scale = None
+            self.per_layer_model_projection = None
+            self.per_layer_projection_norm = None
+
+        # Arrange for shared KVs
+        self.previous_kvs = list(range(len(self.layers)))
+        if config.num_kv_shared_layers > 0:
+            N = len(self.layers)
+            M = N - config.num_kv_shared_layers
+            kvs_by_type = {}
+            for i in range(M):
+                kvs_by_type[self.layers[i].layer_type] = i
+            for j in range(M, N):
+                self.previous_kvs[j] = kvs_by_type[self.layers[j].layer_type]
+
+    def _get_per_layer_inputs(
+        self,
+        input_ids: Optional[mx.array],
+        input_embeddings: Optional[mx.array] = None,
+    ) -> mx.array:
+        if input_ids is None:
+            if input_embeddings is None:
+                raise RuntimeError(
+                    "input_embeddings must be provided when input_ids are omitted."
+                )
+
+            # Split the sequence dimension if this still holds too much
+            # memory. 260k vocab means the distance tensor would be ~1GB
+            # per 2k tokens in bf16.
+            #
+            # If the embedding is quantized we have to dequantize it anyway to
+            # perform the match test.
+            norms_embedding = self.embed_tokens.weight.square().sum(-1)
+            norms_input = input_embeddings.square().sum(-1)
+            distance = _complete_square(
+                norms_embedding,
+                norms_input,
+                self.embed_tokens.as_linear(input_embeddings),
+            )
+
+            # Checks can be added if needed but they necessarily break the GPU
+            # pipelining and force an eval.
+            #
+            #   match_counts = (distance < eps).sum(-1)
+            #
+            input_ids = mx.argmin(distance, -1)
+
+        result = self.embed_tokens_per_layer(input_ids)
+        result = result * self.embed_tokens_per_layer_scale
+        return mx.unflatten(
+            result,
+            -1,
+            (self.config.num_hidden_layers, self.hidden_size_per_layer_input),
+        )
+
+    def _project_per_layer_inputs(
+        self,
+        input_embeddings: mx.array,
+        per_layer_inputs: Optional[mx.array] = None,
+    ) -> mx.array:
+        per_layer_projection = self.per_layer_model_projection(input_embeddings)
+        per_layer_projection = per_layer_projection * self.per_layer_projection_scale
+        per_layer_projection = mx.unflatten(
+            per_layer_projection,
+            -1,
+            (self.config.num_hidden_layers, self.hidden_size_per_layer_input),
+        )
+        per_layer_projection = self.per_layer_projection_norm(per_layer_projection)
+
+        if per_layer_inputs is None:
+            return per_layer_projection
+
+        return (per_layer_projection + per_layer_inputs) * self.per_layer_input_scale
+
+    def _make_masks(self, h, cache):
+        mask = {}
+        masks = []
+        for l, c in zip(self.layers, cache):
+            if l.layer_type not in mask:
+                if l.layer_type == "full_attention":
+                    mask["full_attention"] = create_attention_mask(h, c)
+                elif l.layer_type == "sliding_attention":
+                    mask["sliding_attention"] = create_attention_mask(
+                        h, c, window_size=self.window_size
+                    )
+            masks.append(mask[l.layer_type])
+        return masks
+
+    def __call__(
+        self,
+        inputs: mx.array = None,
+        cache=None,
+        input_embeddings: Optional[mx.array] = None,
+        per_layer_inputs: Optional[mx.array] = None,
+    ):
+        # Make the initial hidden state
+        if input_embeddings is None:
+            input_embeddings = self.embed_tokens(inputs)
+        h = input_embeddings
+        h = h * self.embed_scale
+
+        # Get the extra inputs per layer if we have per layer embeddings
+        if self.hidden_size_per_layer_input:
+            if per_layer_inputs is None:
+                per_layer_inputs = self._get_per_layer_inputs(inputs, input_embeddings)
+            per_layer_inputs = self._project_per_layer_inputs(h, per_layer_inputs)
+        if per_layer_inputs is not None:
+            per_layer_inputs = [
+                per_layer_inputs[:, :, i, :] for i, _ in enumerate(self.layers)
+            ]
+        else:
+            per_layer_inputs = [None] * len(self.layers)
+
+        # Make the kv cache list, be sure to append None for all the shared kv
+        # layers
+        if cache is None:
+            cache = [None] * len(self.layers)
+        else:
+            cache = cache + [None] * (len(self.layers) - len(cache))
+
+        # Apply each layer. We save all intermediate kvs and offset and grab
+        # the previous one for the shared kv layers.
+        masks = self._make_masks(h, cache)
+        intermediates = [(None, None)] * len(self.layers)
+        for idx, (layer, c, mask, prev_idx, per_layer_input) in enumerate(
+            zip(
+                self.layers,
+                cache,
+                masks,
+                self.previous_kvs,
+                per_layer_inputs,
+            )
+        ):
+            kvs, offset = intermediates[prev_idx]
+
+            h, kvs, offset = layer(
+                h,
+                mask,
+                c,
+                per_layer_input=per_layer_input,
+                shared_kv=kvs,
+                offset=offset,
+            )
+
+            intermediates[idx] = (kvs, offset)
+
+        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 = Gemma4TextModel(args)
+        self.final_logit_softcapping = args.final_logit_softcapping
+        self.tie_word_embeddings = args.tie_word_embeddings
+        if not self.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,
+        per_layer_inputs: Optional[mx.array] = None,
+    ):
+        out = self.model(
+            inputs,
+            cache=cache,
+            input_embeddings=input_embeddings,
+            per_layer_inputs=per_layer_inputs,
+        )
+        if self.tie_word_embeddings:
+            out = self.model.embed_tokens.as_linear(out)
+        else:
+            out = self.lm_head(out)
+        if self.final_logit_softcapping is not None:
+            out = logit_softcap(self.final_logit_softcapping, out)
+        return out
+
+    def sanitize(self, weights):
+        sanitized = {}
+        first_kv_shared = self.args.num_hidden_layers - self.args.num_kv_shared_layers
+        for k, v in weights.items():
+            if any(
+                s in k
+                for s in (
+                    "self_attn.rotary_emb",
+                    "input_max",
+                    "input_min",
+                    "output_max",
+                    "output_min",
+                )
+            ):
+                continue
+
+            # KV-shared layers reuse K/V from earlier layers — drop their projections
+            if any(
+                s in k
+                for s in (".self_attn.k_proj", ".self_attn.v_proj", ".self_attn.k_norm")
+            ):
+                try:
+                    layer_idx = int(k.split("layers.")[1].split(".")[0])
+                    if layer_idx >= first_kv_shared:
+                        continue
+                except (IndexError, ValueError):
+                    pass
+
+            if k.endswith(".experts.gate_up_proj"):
+                base = k.removesuffix(".gate_up_proj")
+                gate, up = map(mx.contiguous, mx.split(v, 2, axis=-2))
+                sanitized[f"{base}.switch_glu.gate_proj.weight"] = gate
+                sanitized[f"{base}.switch_glu.up_proj.weight"] = up
+                continue
+
+            if k.endswith(".experts.down_proj"):
+                base = k.removesuffix(".down_proj")
+                sanitized[f"{base}.switch_glu.down_proj.weight"] = v
+                continue
+
+            sanitized[k] = v
+
+        return sanitized
+
+    @property
+    def quant_predicate(self):
+        def predicate(path, _):
+            if path.endswith("router.proj"):
+                return {"group_size": 64, "bits": 8}
+            return True
+
+        return predicate
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    @property
+    def head_dim(self):
+        return self.args.head_dim
+
+    @property
+    def n_kv_heads(self):
+        return self.args.num_key_value_heads
+
+    def make_cache(self):
+        first_kv_shared = self.args.num_hidden_layers - self.args.num_kv_shared_layers
+        caches = []
+        for i in range(first_kv_shared):
+            if self.args.layer_types[i] == "full_attention":
+                caches.append(KVCache())
+            else:
+                caches.append(
+                    RotatingKVCache(
+                        max_size=self.args.sliding_window,
+                        keep=0,
+                    )
+                )
+        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,7 +7,9 @@ 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
@@ -122,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
 
 
@@ -205,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
 
 
@@ -252,10 +262,6 @@ class LanguageModel(PipelineMixin, 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__(
@@ -286,7 +292,8 @@ class LanguageModel(PipelineMixin, nn.Module):
                 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)
 
@@ -329,6 +336,61 @@ 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.pipeline_layers

mlx_lm/models/glm4_moe_lite.py

@@ -0,0 +1,531 @@
+# 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 .mla import MultiLinear
+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
+    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
+    quantization: Optional[Dict[str, Any]] = None
+
+
+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
+        rope_params = config.rope_scaling
+        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)
+        head_dim = self.qk_nope_head_dim + self.v_head_dim
+        self.embed_q = MultiLinear(
+            self.qk_nope_head_dim, self.kv_lora_rank, self.num_heads
+        )
+        self.unembed_out = MultiLinear(
+            self.kv_lora_rank, self.v_head_dim, self.num_heads
+        )
+
+        self.o_proj = nn.Linear(
+            self.num_heads * self.v_head_dim,
+            self.hidden_size,
+            bias=config.attention_bias,
+        )
+
+        if rope_params is not None:
+            mscale_all_dim = rope_params.get("mscale_all_dim", 0)
+            if mscale_all_dim:
+                scaling_factor = rope_params["factor"]
+                if scaling_factor > 1:
+                    s = 0.1 * mscale_all_dim * math.log(scaling_factor) + 1.0
+                    self.scale = self.scale * s * s
+
+        self.rope = initialize_rope(
+            dims=self.qk_rope_head_dim,
+            base=self.rope_theta,
+            traditional=True,
+            max_position_embeddings=self.max_position_embeddings,
+            scaling_config=rope_params,
+        )
+
+    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_latent = self.kv_a_layernorm(compressed_kv)
+
+        offset = cache.offset if cache is not None else 0
+        q_pe = self.rope(q_pe, offset)
+        k_pe = self.rope(k_pe, offset)
+
+        kv_latent = mx.expand_dims(kv_latent, axis=1)
+
+        if cache is not None:
+            kv_latent, k_pe = cache.update_and_fetch(kv_latent, k_pe)
+
+        pe_scores = (q_pe * self.scale) @ k_pe.swapaxes(-1, -2)
+        if mask is not None:
+            pe_scores = mx.where(
+                mask,
+                pe_scores,
+                mx.array(mx.finfo(pe_scores.dtype).min, pe_scores.dtype),
+            )
+
+        if L == 1:
+            q_nope = self.embed_q(q_nope)
+            k = v = kv_latent
+        else:
+            k = self.embed_q(kv_latent, transpose=False)
+            v = self.unembed_out(kv_latent)
+        output = scaled_dot_product_attention(
+            q_nope, k, v, cache=cache, scale=self.scale, mask=pe_scores
+        )
+        if L == 1:
+            output = self.unembed_out(output)
+
+        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 + 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,))
+        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)
+        use_moe = (
+            config.n_routed_experts is not None
+            and layer_idx >= config.first_k_dense_replace
+            and layer_idx % config.moe_layer_freq == 0
+        )
+        self.mlp = Glm4MoeLiteMoE(config) if use_moe 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], 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 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):
+        def is_mpt_layer(key):
+            subkeys = key.split(".")
+            if len(subkeys) < 3:
+                return False
+            if (
+                subkeys[1] == "layers"
+                and int(subkeys[2]) >= self.args.num_hidden_layers
+            ):
+                return True
+            return False
+
+        new_weights = {}
+        for k, v in weights.items():
+            if is_mpt_layer(k):
+                continue
+            else:
+                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)
+            prefix = f"model.layers.{l}.self_attn"
+            if f"{prefix}.kv_b_proj.weight" in weights:
+                layer = self.layers[l].self_attn.embed_q
+                quantized = f"{prefix}.kv_b_proj.scales" in weights
+                v = weights.pop(f"{prefix}.kv_b_proj.weight")
+                head_dim = self.args.qk_nope_head_dim + self.args.v_head_dim
+
+                if quantized:
+                    dims = self.args.kv_lora_rank
+                    scales = weights.pop(f"{prefix}.kv_b_proj.scales")
+                    biases = weights.pop(f"{prefix}.kv_b_proj.biases")
+                    # Try to infer bits and group size
+                    bits = (v.shape[-1] * 32) // dims
+                    group_size = dims // scales.shape[-1]
+                    v = mx.dequantize(
+                        v, scales, biases, bits=bits, group_size=group_size
+                    )
+                num_heads = self.args.num_attention_heads
+                v = v.reshape(num_heads, head_dim, -1)
+                wk = mx.contiguous(
+                    v[:, : self.args.qk_nope_head_dim, :].swapaxes(-1, -2)
+                )
+                wv = mx.contiguous(v[:, self.args.qk_nope_head_dim :, :])
+                if quantized:
+                    wk, wk_scales, wk_biases = mx.quantize(
+                        wk, bits=bits, group_size=group_size
+                    )
+                    wv, wv_scales, wv_biases = mx.quantize(
+                        wv, bits=bits, group_size=group_size
+                    )
+                    weights[f"{prefix}.embed_q.scales"] = wk_scales
+                    weights[f"{prefix}.unembed_out.scales"] = wv_scales
+                    weights[f"{prefix}.embed_q.biases"] = wk_biases
+                    weights[f"{prefix}.unembed_out.biases"] = wv_biases
+                weights[f"{prefix}.embed_q.weight"] = wk
+                weights[f"{prefix}.unembed_out.weight"] = wv
+
+        return weights
+
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        group = group or mx.distributed.init()
+        rank = group.rank()
+        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.num_heads //= N
+            num_heads = layer.self_attn.num_heads
+            sh = rank * num_heads
+            eh = sh + num_heads
+
+            def shard_heads(w):
+                return w[sh:eh]
+
+            layer.self_attn.embed_q.apply(shard_heads)
+            layer.self_attn.unembed_out.apply(shard_heads)
+
+            layer.self_attn.o_proj = shard_linear(
+                layer.self_attn.o_proj, "sharded-to-all", group=group
+            )
+
+            # 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/glm_moe_dsa.py

@@ -0,0 +1,53 @@
+# Copyright © 2025 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+
+from .base import BaseModelArgs
+from .deepseek_v32 import Model as DSV32Model
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    vocab_size: int
+    hidden_size: int
+    index_head_dim: int
+    index_n_heads: int
+    index_topk: 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: float
+    kv_lora_rank: int
+    q_lora_rank: int
+    qk_rope_head_dim: int
+    v_head_dim: int
+    qk_nope_head_dim: int
+    topk_method: str
+    scoring_func: str
+    norm_topk_prob: bool
+    n_group: int
+    topk_group: int
+    num_experts_per_tok: int
+    moe_layer_freq: int
+    first_k_dense_replace: int
+    max_position_embeddings: int
+    rms_norm_eps: float
+    rope_parameters: Dict
+    attention_bias: bool
+    rope_scaling: Dict = None
+    rope_theta: Optional[float] = None
+
+    def __post_init__(self):
+        self.rope_scaling = self.rope_parameters
+        self.rope_theta = self.rope_parameters["rope_theta"]
+
+
+class Model(DSV32Model):
+    def __init__(self, config: ModelArgs):
+        super().__init__(config)

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,13 +6,14 @@ 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,
     create_ssm_mask,
     scaled_dot_product_attention,
 )
-from .cache import KVCache, MambaCache
+from .cache import ArraysCache, KVCache
 from .rope_utils import initialize_rope
 from .ssm import ssm_update
 from .switch_layers import SwitchGLU
@@ -75,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)
 
 
@@ -119,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[ArraysCache],
+        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)
@@ -144,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[ArraysCache],
+        mask: Optional[mx.array],
     ) -> mx.array:
         batch_size, seq_len, _ = hidden_states.shape
 
@@ -154,27 +170,34 @@ 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,
-        cache: Optional[MambaCache] = None,
+        mask: Optional[mx.array],
+        cache: Optional[ArraysCache] = None,
     ) -> mx.array:
 
         projected = self.in_proj(hidden_states)
@@ -184,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,
             [
@@ -197,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)
 
@@ -320,7 +338,7 @@ 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):
@@ -335,7 +353,7 @@ class GraniteMoeHybridMLP(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 GraniteMoeHybridLayer(nn.Module):
@@ -478,7 +496,7 @@ class Model(nn.Module):
         caches = []
         for layer in self.layers:
             if layer.layer_type == "mamba":
-                caches.append(MambaCache())
+                caches.append(ArraysCache(size=2))
             elif layer.layer_type == "attention":
                 caches.append(KVCache())
         return caches

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

@@ -7,13 +7,14 @@ 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 .cache import ArraysCache, KVCache
 from .switch_layers import SwitchGLU
 
 
@@ -65,7 +66,7 @@ class JambaMLP(nn.Module):
         self.down_proj = nn.Linear(args.intermediate_size, args.hidden_size, bias=False)
 
     def __call__(self, x: mx.array) -> mx.array:
-        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class JambaAttention(nn.Module):
@@ -205,7 +206,7 @@ class JambaMambaMixer(nn.Module):
         x = nn.silu(conv_out)
         A = -mx.exp(self.A_log)
         y, ssm_state = self.ssm_step(x, A, ssm_state)
-        z = self.out_proj(nn.silu(z) * y)
+        z = self.out_proj(swiglu(z, y))
         return z, (conv_state, ssm_state)
 
     def __call__(self, x, cache):
@@ -340,7 +341,7 @@ class Model(nn.Module):
             if layer.is_attn:
                 caches.append(KVCache())
             else:
-                caches.append(MambaCache())
+                caches.append(ArraysCache(size=2))
         return caches
 
     def sanitize(self, weights):

mlx_lm/models/kimi_k25.py

@@ -0,0 +1,83 @@
+# Copyright © 2026 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Optional, Union
+
+import mlx.core as mx
+import mlx.nn as nn
+from mlx.utils import tree_flatten, tree_unflatten
+
+from .base import BaseModelArgs
+from .deepseek_v3 import DeepseekV3Model
+from .deepseek_v3 import Model as DeepseekV3LM
+from .deepseek_v3 import ModelArgs as TextConfig
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    text_config: Union[TextConfig, dict]
+    model_type: str = "kimi_k25"
+
+    def __post_init__(self):
+        if isinstance(self.text_config, dict):
+            self.text_config = TextConfig.from_dict(self.text_config)
+
+
+class LanguageModel(nn.Module):
+    def __init__(self, config: TextConfig):
+        super().__init__()
+        self.args = config
+        self.model = DeepseekV3Model(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ):
+        out = self.model(inputs, cache)
+        return self.lm_head(out)
+
+
+class Model(nn.Module):
+    def __init__(self, config: ModelArgs):
+        super().__init__()
+        self.args = config
+        self.model_type = config.model_type
+        self.language_model = LanguageModel(config.text_config)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ):
+        return self.language_model(inputs, cache)
+
+    def sanitize(self, weights):
+        weights = tree_unflatten(list(weights.items()))
+        weights.pop("vision_tower", None)
+        weights.pop("vision_model", None)
+        weights.pop("multi_modal_projector", None)
+        weights.pop("mm_projector", None)
+        lm_weights = dict(tree_flatten(weights["language_model"]))
+        lm_weights = DeepseekV3LM.sanitize(self.language_model, lm_weights)
+        weights["language_model"] = tree_unflatten(list(lm_weights.items()))
+        return dict(tree_flatten(weights))
+
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        DeepseekV3LM.shard(self.language_model, group)
+
+    @property
+    def model(self):
+        return self.language_model.model
+
+    @property
+    def layers(self):
+        return self.language_model.model.pipeline_layers
+
+    @property
+    def cast_predicate(self):
+        def predicate(k):
+            return "e_score_correction_bias" not in k
+
+        return predicate

mlx_lm/models/kimi_linear.py

@@ -6,15 +6,16 @@ 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 .cache import ArraysCache, KVCache
 from .gated_delta import gated_delta_update
-from .rope_utils import initialize_rope
+from .mla import MultiLinear
 from .switch_layers import SwitchGLU
 
 
@@ -68,7 +69,7 @@ class KimiMLP(nn.Module):
         self.down_proj = nn.Linear(hidden, dim, bias=False)
 
     def __call__(self, x: mx.array) -> mx.array:
-        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 @mx.compile
@@ -164,6 +165,7 @@ class KimiMLAAttention(nn.Module):
         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.kv_lora_rank = args.kv_lora_rank
         self.scale = self.q_head_dim**-0.5
 
         hidden = args.hidden_size
@@ -174,23 +176,14 @@ class KimiMLAAttention(nn.Module):
             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.embed_q = MultiLinear(
+            self.qk_nope_head_dim, args.kv_lora_rank, self.num_heads
+        )
+        self.unembed_out = MultiLinear(
+            args.kv_lora_rank, self.v_head_dim, self.num_heads
         )
         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,
@@ -198,51 +191,45 @@ class KimiMLAAttention(nn.Module):
         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)
+        q = self.q_proj(x).reshape(B, L, self.num_heads, self.q_head_dim)
+        q = q.transpose(0, 2, 1, 3)
+        q_nope, q_pe = mx.split(q, [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)
+        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_latent = self.kv_a_layernorm(compressed_kv)
 
-        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)
+        kv_latent = mx.expand_dims(kv_latent, axis=1)
 
         if cache is not None:
-            queries = self.rope(queries, offset=cache.offset)
-            keys = self.rope(keys, offset=cache.offset)
-            keys, values = cache.update_and_fetch(keys, values)
-        else:
-            queries = self.rope(queries)
-            keys = self.rope(keys)
+            kv_latent, k_pe = cache.update_and_fetch(kv_latent, k_pe)
 
-        out = scaled_dot_product_attention(
-            queries,
-            keys,
-            values,
-            cache,
-            scale=self.scale,
-            mask=mask,
+        pe_scores = (q_pe * self.scale) @ k_pe.swapaxes(-1, -2)
+        if mask is not None:
+            pe_scores = mx.where(
+                mask,
+                pe_scores,
+                mx.array(mx.finfo(pe_scores.dtype).min, pe_scores.dtype),
+            )
+
+        if L == 1:
+            q_nope = self.embed_q(q_nope)
+            k = v = kv_latent
+        else:
+            k = self.embed_q(kv_latent, transpose=False)
+            v = self.unembed_out(kv_latent)
+
+        output = scaled_dot_product_attention(
+            q_nope, k, v, cache=cache, scale=self.scale, mask=pe_scores
         )
-        out = out.transpose(0, 2, 1, 3).reshape(B, L, -1)
-        return self.o_proj(out)
+
+        if L == 1:
+            output = self.unembed_out(output)
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
 
 
 class ShortConv1d(nn.Module):
@@ -259,18 +246,30 @@ class ShortConv1d(nn.Module):
         )
 
     def __call__(
-        self, x: mx.array, cache: Optional[mx.array]
+        self,
+        x: mx.array,
+        state: Optional[mx.array],
+        mask: Optional[mx.array],
+        lengths: Optional[mx.array],
     ) -> Tuple[mx.array, mx.array]:
-        if cache is None:
-            pad = mx.zeros(
+        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
             )
-        else:
-            pad = cache
-        conv_input = mx.concatenate([pad, x], axis=1)
+        conv_input = mx.concatenate([state, x], axis=1)
         out = nn.silu(self.conv(conv_input))
-        new_cache = conv_input[:, -self.kernel_size + 1 :, :]
-        return out, new_cache
+        n_keep = self.kernel_size - 1
+        if lengths is not None:
+            ends = mx.clip(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 = mx.contiguous(conv_input[:, -n_keep:, :])
+
+        return out, new_state
 
 
 class KimiDeltaAttention(nn.Module):
@@ -322,37 +321,37 @@ class KimiDeltaAttention(nn.Module):
         dtype = x.dtype
 
         if cache is not None:
-            conv_state, ssm_state = cache
+            q_state, k_state, v_state, ssm_state = cache
+            lengths = cache.lengths
         else:
-            conv_state = None
+            q_state = None
+            k_state = None
+            v_state = None
             ssm_state = None
+            lengths = None
 
-        if conv_state is None:
+        if q_state is None:
             s = mx.zeros((B, self.conv_kernel - 1, self.projection_dim), dtype=dtype)
             q_state = s
             k_state = s
             v_state = s
-        else:
-            q_state, k_state, v_state = conv_state
 
-        q_conv, q_state = self.q_conv(self.q_proj(x), q_state)
-        k_conv, k_state = self.k_conv(self.k_proj(x), k_state)
-        v_conv, v_state = self.v_conv(self.v_proj(x), v_state)
+        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)
+            cache[0] = q_state
+            cache[1] = k_state
+            cache[2] = 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
+        inv_scale = self.scale
+        q = (inv_scale**2) * mx.fast.rms_norm(q, None, 1e-6)
+        k = inv_scale * mx.fast.rms_norm(k, None, 1e-6)
 
         a_logits = self.f_b_proj(self.f_a_proj(x)).reshape(
             B, T, self.num_heads, self.head_dim
@@ -373,7 +372,8 @@ class KimiDeltaAttention(nn.Module):
         )
 
         if cache is not None:
-            cache[1] = ssm_state
+            cache[3] = ssm_state
+            cache.advance(T)
 
         gate = self.g_b_proj(self.g_a_proj(x)).reshape(
             B, T, self.num_heads, self.head_dim
@@ -446,7 +446,7 @@ class KimiLinearModel(nn.Module):
             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])
+        attn_mask = create_attention_mask(h, cache[self.attn_idx], return_array=True)
 
         for layer, layer_cache in zip(self.layers, cache):
             mask = ssm_mask if layer.is_linear else attn_mask
@@ -484,7 +484,7 @@ class Model(nn.Module):
         caches: List[Any] = []
         for layer in self.layers:
             if layer.is_linear:
-                caches.append(MambaCache())
+                caches.append(ArraysCache(size=4))
             else:
                 caches.append(KVCache())
         return caches
@@ -552,6 +552,42 @@ class Model(nn.Module):
                     if weights[dt_key].ndim > 1:
                         weights[dt_key] = mx.reshape(weights[dt_key], (-1,))
 
+            attn_prefix = f"{prefix}.self_attn"
+            kv_b_key = f"{attn_prefix}.kv_b_proj.weight"
+            if kv_b_key in weights:
+                qk_nope = self.args.qk_nope_head_dim or self.args.head_dim
+                v_head = self.args.v_head_dim or self.args.head_dim
+                head_dim = qk_nope + v_head
+                num_heads = self.args.num_attention_heads
+
+                quantized = f"{attn_prefix}.kv_b_proj.scales" in weights
+                v = weights.pop(kv_b_key)
+
+                if quantized:
+                    dims = self.args.kv_lora_rank
+                    scales = weights.pop(f"{attn_prefix}.kv_b_proj.scales")
+                    biases = weights.pop(f"{attn_prefix}.kv_b_proj.biases")
+                    bits = (v.shape[-1] * 32) // dims
+                    group_size = dims // scales.shape[-1]
+                    v = mx.dequantize(
+                        v, scales, biases, bits=bits, group_size=group_size
+                    )
+
+                v = v.reshape(num_heads, head_dim, -1)
+                wk = mx.contiguous(v[:, :qk_nope, :].swapaxes(-1, -2))
+                wv = mx.contiguous(v[:, qk_nope:, :])
+
+                if quantized:
+                    wk, wk_s, wk_b = mx.quantize(wk, bits=bits, group_size=group_size)
+                    wv, wv_s, wv_b = mx.quantize(wv, bits=bits, group_size=group_size)
+                    weights[f"{attn_prefix}.embed_q.scales"] = wk_s
+                    weights[f"{attn_prefix}.embed_q.biases"] = wk_b
+                    weights[f"{attn_prefix}.unembed_out.scales"] = wv_s
+                    weights[f"{attn_prefix}.unembed_out.biases"] = wv_b
+
+                weights[f"{attn_prefix}.embed_q.weight"] = wk
+                weights[f"{attn_prefix}.unembed_out.weight"] = wv
+
         return weights
 
     @property

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,11 +32,14 @@ class ModelArgs(BaseModelArgs):
     block_multiple_of: int
     block_ffn_dim_multiplier: float
     block_auto_adjust_ff_dim: bool
-    rope_theta: float
+    rope_theta: float = 1000000.0
+    rope_parameters: Optional[dict] = None
     full_attn_idxs: Optional[List[int]] = None
     layer_types: Optional[List[str]] = None
 
     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"]
         if self.num_key_value_heads is None:
             self.num_key_value_heads = self.num_attention_heads
         if self.full_attn_idxs is None:
@@ -138,17 +142,28 @@ class ShortConv(nn.Module):
         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
@@ -176,7 +191,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

@@ -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,
@@ -34,11 +35,14 @@ class ModelArgs(BaseModelArgs):
     norm_eps: float
     conv_bias: bool
     conv_L_cache: int
-    rope_theta: float
+    rope_theta: float = 1000000.0
+    rope_parameters: Optional[dict] = None
     full_attn_idxs: Optional[List[int]] = None
     layer_types: Optional[List[str]] = None
 
     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"]
         if self.full_attn_idxs is None:
             self.full_attn_idxs = [
                 i
@@ -139,17 +143,28 @@ class ShortConv(nn.Module):
         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
@@ -168,7 +183,7 @@ class MLP(nn.Module):
         self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
 
     def __call__(self, x) -> mx.array:
-        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class Lfm2MoeSparseMoeBlock(nn.Module):

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

@@ -5,7 +5,9 @@ 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
@@ -116,7 +118,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):
@@ -226,6 +228,37 @@ 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

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,13 @@ 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 .mla import MultiLinear
+from .rope_utils import initialize_rope
 from .switch_layers import SwitchGLU
 
 
@@ -38,6 +42,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):
@@ -76,10 +81,11 @@ class LongcatFlashMLA(nn.Module):
             bias=args.attention_bias,
         )
         self.kv_a_layernorm = nn.RMSNorm(self.kv_lora_rank)
-        self.kv_b_proj = nn.Linear(
-            self.kv_lora_rank,
-            self.num_attention_heads * (self.qk_nope_head_dim + args.v_head_dim),
-            bias=False,
+        self.embed_q = MultiLinear(
+            self.qk_nope_head_dim, self.kv_lora_rank, self.num_attention_heads
+        )
+        self.unembed_out = MultiLinear(
+            self.kv_lora_rank, self.v_head_dim, self.num_attention_heads
         )
 
         self.o_proj = nn.Linear(
@@ -93,8 +99,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__(
@@ -106,56 +124,59 @@ class LongcatFlashMLA(nn.Module):
         B, L, _ = x.shape
 
         if self.q_lora_rank is None:
-            q_states = self.q_proj(x)
+            q = self.q_proj(x)
         else:
-            q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x)))
+            q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(x)))
 
-        q_states = q_states.reshape(B, L, -1, self.qk_head_dim).transpose(0, 2, 1, 3)
-
-        if self.mla_scale_q_lora is not None:
-            q_states = q_states * self.mla_scale_q_lora
-
-        q_pass, q_rot = mx.split(q_states, [self.qk_nope_head_dim], axis=-1)
-
-        compressed_kv = self.kv_a_proj_with_mqa(x)
-        k_pass, k_rot = mx.split(compressed_kv, [self.kv_lora_rank], axis=-1)
-        k_pass = self.kv_a_layernorm(k_pass)
-
-        if self.mla_scale_kv_lora is not None:
-            k_pass = k_pass * self.mla_scale_kv_lora
-
-        key_shape = (B, L, -1, self.qk_nope_head_dim + self.v_head_dim)
-        k_pass = self.kv_b_proj(k_pass).reshape(*key_shape).transpose(0, 2, 1, 3)
-        k_pass, value_states = mx.split(k_pass, [self.qk_nope_head_dim], axis=-1)
-
-        k_rot = k_rot.reshape(B, 1, L, self.qk_rope_head_dim)
-
-        if cache is not None:
-            q_rot = self.rope(q_rot, cache.offset)
-            k_rot = self.rope(k_rot, cache.offset)
-        else:
-            q_rot = self.rope(q_rot)
-            k_rot = self.rope(k_rot)
-
-        k_rot = mx.broadcast_to(k_rot, (*k_pass.shape[:-1], k_rot.shape[-1]))
-
-        query_states = mx.concatenate([q_pass, q_rot], axis=-1)
-        key_states = mx.concatenate([k_pass, k_rot], axis=-1)
-
-        if cache is not None:
-            key_states, value_states = cache.update_and_fetch(key_states, value_states)
-
-        attn_output = scaled_dot_product_attention(
-            query_states,
-            key_states,
-            value_states,
-            cache=cache,
-            scale=self.scale,
-            mask=mask,
+        q = q.reshape(B, L, self.num_attention_heads, self.qk_head_dim).transpose(
+            0, 2, 1, 3
         )
 
-        attn_output = attn_output.transpose(0, 2, 1, 3).reshape(B, L, -1)
-        return self.o_proj(attn_output)
+        if self.mla_scale_q_lora is not None:
+            q = q * self.mla_scale_q_lora
+
+        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_latent = self.kv_a_layernorm(compressed_kv)
+
+        if self.mla_scale_kv_lora is not None:
+            kv_latent = kv_latent * self.mla_scale_kv_lora
+
+        offset = cache.offset if cache is not None else 0
+        q_pe = self.rope(q_pe, offset)
+        k_pe = self.rope(k_pe, offset)
+
+        kv_latent = mx.expand_dims(kv_latent, axis=1)
+
+        if cache is not None:
+            kv_latent, k_pe = cache.update_and_fetch(kv_latent, k_pe)
+
+        pe_scores = (q_pe * self.scale) @ k_pe.swapaxes(-1, -2)
+        if mask is not None:
+            pe_scores = mx.where(
+                mask,
+                pe_scores,
+                mx.array(mx.finfo(pe_scores.dtype).min, pe_scores.dtype),
+            )
+
+        if L == 1:
+            q_nope = self.embed_q(q_nope)
+            k = v = kv_latent
+        else:
+            k = self.embed_q(kv_latent, transpose=False)
+            v = self.unembed_out(kv_latent)
+
+        output = scaled_dot_product_attention(
+            q_nope, k, v, cache=cache, scale=self.scale, mask=pe_scores
+        )
+        if L == 1:
+            output = self.unembed_out(output)
+
+        output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.o_proj(output)
 
 
 class LongcatFlashMLP(nn.Module):
@@ -168,7 +189,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 +244,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)
 
@@ -236,14 +260,20 @@ class LongcatFlashMoE(nn.Module):
         regular_outputs = self.switch_mlp(hidden_states, topk_indices)
 
         weighted_outputs = regular_outputs * regular_weights[..., None]
-
-        # Add identity expert contribution if needed
-        assert self.zero_expert_type == "identity"
-        identity_weights = mx.where(mask, topk_weights, 0.0)
-        identity_outputs = hidden_states[..., None, :] * identity_weights[..., None]
-        weighted_outputs = weighted_outputs + identity_outputs
-
         final_output = mx.sum(weighted_outputs, axis=-2)
+
+        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
 
 
@@ -314,7 +344,7 @@ class LongcatFlashModel(nn.Module):
         if cache is None:
             cache = [(None, None)] * self.num_layers
 
-        mask = create_attention_mask(h, cache[0][0])
+        mask = create_attention_mask(h, cache[0][0], return_array=True)
 
         for layer, c in zip(self.layers, cache):
             h = layer(h, mask, cache=c)
@@ -370,6 +400,47 @@ class Model(nn.Module):
                         ]
                         weights[f"{prefix}.mlp.switch_mlp.{m}.{k}"] = mx.stack(to_join)
 
+        for l in range(self.args.num_layers):
+            for i in range(2):
+                prefix = f"model.layers.{l}.self_attn.{i}"
+                kv_b_key = f"{prefix}.kv_b_proj.weight"
+                if kv_b_key in weights:
+                    num_heads = self.args.num_attention_heads
+                    head_dim = self.args.qk_nope_head_dim + self.args.v_head_dim
+                    quantized = f"{prefix}.kv_b_proj.scales" in weights
+                    v = weights.pop(kv_b_key)
+
+                    if quantized:
+                        dims = self.args.kv_lora_rank
+                        scales = weights.pop(f"{prefix}.kv_b_proj.scales")
+                        biases = weights.pop(f"{prefix}.kv_b_proj.biases")
+                        bits = (v.shape[-1] * 32) // dims
+                        group_size = dims // scales.shape[-1]
+                        v = mx.dequantize(
+                            v, scales, biases, bits=bits, group_size=group_size
+                        )
+
+                    v = v.reshape(num_heads, head_dim, -1)
+                    wk = mx.contiguous(
+                        v[:, : self.args.qk_nope_head_dim, :].swapaxes(-1, -2)
+                    )
+                    wv = mx.contiguous(v[:, self.args.qk_nope_head_dim :, :])
+
+                    if quantized:
+                        wk, wk_s, wk_b = mx.quantize(
+                            wk, bits=bits, group_size=group_size
+                        )
+                        wv, wv_s, wv_b = mx.quantize(
+                            wv, bits=bits, group_size=group_size
+                        )
+                        weights[f"{prefix}.embed_q.scales"] = wk_s
+                        weights[f"{prefix}.embed_q.biases"] = wk_b
+                        weights[f"{prefix}.unembed_out.scales"] = wv_s
+                        weights[f"{prefix}.unembed_out.biases"] = wv_b
+
+                    weights[f"{prefix}.embed_q.weight"] = wk
+                    weights[f"{prefix}.unembed_out.weight"] = wv
+
         new_weights = {}
         for k, v in weights.items():
             if k.startswith("model.mtp"):
@@ -379,3 +450,44 @@ 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()
+        rank = group.rank()
+
+        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.o_proj = shard_linear(attn.o_proj, "sharded-to-all", group=group)
+                attn.num_attention_heads //= N
+                num_heads = attn.num_attention_heads
+                sh = rank * num_heads
+                eh = sh + num_heads
+
+                def shard_heads(w):
+                    return w[sh:eh]
+
+                attn.embed_q.apply(shard_heads)
+                attn.unembed_out.apply(shard_heads)
+
+            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/longcat_flash_ngram.py

@@ -0,0 +1,214 @@
+# Copyright © 2026 Apple Inc.
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from .base import BaseModelArgs, create_attention_mask
+from .cache import ArraysCache, CacheList, KVCache
+from .longcat_flash import LongcatFlashDecoderLayer
+from .longcat_flash import Model as LongcatFlashLM
+
+
+@dataclass
+class ModelArgs(BaseModelArgs):
+    model_type: str
+    hidden_size: int
+    ffn_hidden_size: int
+    moe_topk: int
+    expert_ffn_hidden_size: int
+    n_routed_experts: int
+    zero_expert_num: int
+    num_layers: int
+    vocab_size: int
+    max_position_embeddings: int
+    num_attention_heads: int
+    kv_lora_rank: int
+    q_lora_rank: int
+    qk_rope_head_dim: int
+    qk_nope_head_dim: int
+    v_head_dim: int
+    routed_scaling_factor: float
+    rms_norm_eps: float
+    rope_theta: float
+    mla_scale_q_lora: bool
+    mla_scale_kv_lora: bool
+    attention_bias: bool = False
+    zero_expert_type: str = "identity"
+    ngram_vocab_size_ratio: int = 78
+    emb_neighbor_num: int = 4
+    emb_split_num: int = 4
+    norm_topk_prob: bool = False
+    router_bias: bool = False
+    rope_scaling: Optional[Dict] = None
+
+
+class NgramEmbedding(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.vocab_size = args.vocab_size
+        self.hidden_size = args.hidden_size
+        self.m = args.ngram_vocab_size_ratio * args.vocab_size
+        self.k = args.emb_split_num
+        self.n = args.emb_neighbor_num
+
+        self.word_embeddings = nn.Embedding(args.vocab_size, args.hidden_size)
+
+        num_embedders = self.k * (self.n - 1)
+        emb_dim = args.hidden_size // num_embedders
+
+        self.embedders = []
+        self.post_projs = []
+        for i in range(num_embedders):
+            emb_vocab_size = int(self.m + i * 2 + 1)
+            self.embedders.append(nn.Embedding(emb_vocab_size, emb_dim))
+            self.post_projs.append(nn.Linear(emb_dim, args.hidden_size, bias=False))
+        self._compute_vocab_mods()
+
+    def _compute_vocab_mods(self):
+        vocab_mods = {}
+        for i in range(2, self.n + 1):
+            for j in range(self.k):
+                index = (i - 2) * self.k + j
+                emb_vocab_dim = int(self.m + index * 2 + 1)
+                mods = []
+                power_mod = 1
+                for _ in range(i - 1):
+                    power_mod = (power_mod * self.vocab_size) % emb_vocab_dim
+                    mods.append(power_mod)
+                vocab_mods[(i, j)] = mods
+        self._vocab_mods = vocab_mods
+
+    def _shift_right(self, x: mx.array, n: int) -> mx.array:
+        if n <= 0:
+            return x
+        batch_size, seq_len = x.shape
+        if seq_len <= n:
+            return mx.zeros_like(x)
+        return mx.concatenate(
+            [mx.zeros((batch_size, n), dtype=x.dtype), x[..., :-n]], axis=-1
+        )
+
+    def _get_ngram_ids(
+        self,
+        input_ids: mx.array,
+        shifted_ids: Dict[int, mx.array],
+        vocab_mods: List[int],
+        ngram: int,
+    ) -> mx.array:
+        ngram_ids = input_ids
+        for k in range(2, ngram + 1):
+            ngram_ids = ngram_ids + shifted_ids[k] * vocab_mods[k - 2]
+        return ngram_ids
+
+    def __call__(
+        self,
+        input_ids: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        seq_len = input_ids.shape[-1]
+
+        input_ids = input_ids.astype(mx.int64)
+        if cache is not None:
+            context = cache[0]
+            if context is None:
+                context = input_ids
+            else:
+                context = mx.concatenate([context, input_ids], axis=-1)
+            cache[0] = context[..., max(0, context.shape[-1] - self.n + 1) :]
+        else:
+            context = input_ids
+
+        x = self.word_embeddings(input_ids)
+        vocab_mods = self._vocab_mods
+
+        shifted_ids = {}
+        for i in range(2, self.n + 1):
+            shifted_ids[i] = self._shift_right(context, i - 1)
+
+        for i in range(2, self.n + 1):
+            for j in range(self.k):
+                index = (i - 2) * self.k + j
+                emb_vocab_dim = int(self.m + index * 2 + 1)
+                ngram_ids = self._get_ngram_ids(
+                    context, shifted_ids, vocab_mods[(i, j)], ngram=i
+                )
+                new_ids = (ngram_ids % emb_vocab_dim)[..., -seq_len:]
+                x_ngram = self.embedders[index](new_ids)
+                x_proj = self.post_projs[index](x_ngram)
+                x = x + x_proj
+
+        return x / (1 + self.k * (self.n - 1))
+
+
+class LongcatFlashNgramModel(nn.Module):
+    def __init__(self, args: ModelArgs):
+        super().__init__()
+        self.num_layers = args.num_layers
+        self.ngram_embeddings = NgramEmbedding(args)
+        self.layers = [LongcatFlashDecoderLayer(args) for _ in range(args.num_layers)]
+        self.norm = nn.RMSNorm(args.hidden_size, args.rms_norm_eps)
+
+    def __call__(
+        self,
+        input_ids: mx.array,
+        cache: Optional[Any] = None,
+    ) -> mx.array:
+        if cache is None:
+            cache = [None] + [(None, None)] * self.num_layers
+
+        h = self.ngram_embeddings(input_ids, cache=cache[0])
+
+        mask = create_attention_mask(h, cache[1][0], return_array=True)
+
+        for layer, c in zip(self.layers, cache[1:]):
+            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 = LongcatFlashNgramModel(args)
+        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
+
+    def __call__(
+        self,
+        inputs: mx.array,
+        cache: Optional[Any] = None,
+    ):
+        out = self.model(inputs, cache)
+        return self.lm_head(out)
+
+    @property
+    def layers(self):
+        return self.model.layers
+
+    @property
+    def quant_predicate(self):
+        return LongcatFlashLM.quant_predicate.fget(self)
+
+    @property
+    def cast_predicate(self):
+        return LongcatFlashLM.cast_predicate.fget(self)
+
+    def sanitize(self, weights):
+        weights = LongcatFlashLM.sanitize(self, weights)
+        if "model.embed_tokens.weight" in weights:
+            weights["model.ngram_embeddings.word_embeddings.weight"] = weights.pop(
+                "model.embed_tokens.weight"
+            )
+        return weights
+
+    def make_cache(self):
+        return [ArraysCache(size=1)] + [
+            CacheList(KVCache(), KVCache()) for _ in self.model.layers
+        ]
+
+    def shard(self, group: Optional[mx.distributed.Group] = None):
+        LongcatFlashLM.shard(self, group)

mlx_lm/models/mamba.py

@@ -6,8 +6,9 @@ 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
+from .cache import ArraysCache
 
 
 @dataclass
@@ -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):
@@ -152,7 +153,7 @@ class MambaBlock(nn.Module):
             x, conv_cache, state_cache
         )
 
-        if isinstance(cache, MambaCache):
+        if isinstance(cache, ArraysCache):
             cache[0] = new_conv_cache
             cache[1] = new_state_cache
 
@@ -207,7 +208,7 @@ class Model(nn.Module):
         return logits
 
     def make_cache(self):
-        return [MambaCache() for _ in range(len(self.layers))]
+        return [ArraysCache(size=2) for _ in range(len(self.layers))]
 
     @property
     def layers(self):

mlx_lm/models/mamba2.py

@@ -7,8 +7,9 @@ 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 .cache import ArraysCache
 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[ArraysCache],
+        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[ArraysCache],
+        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,14 +159,17 @@ 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,
         hidden_states: mx.array,
         mask: Optional[mx.array],
-        cache: Optional[MambaCache] = None,
+        cache: Optional[ArraysCache] = None,
     ) -> mx.array:
         projected = self.in_proj(hidden_states)
         gate, conv_input, dt = mx.split(
@@ -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)
 
@@ -181,7 +200,7 @@ class ResidualBlock(nn.Module):
         self.norm = nn.RMSNorm(args.hidden_size)
 
     def __call__(
-        self, x: mx.array, mask: Optional[mx.array], cache: Optional[MambaCache] = None
+        self, x: mx.array, mask: Optional[mx.array], cache: Optional[ArraysCache] = None
     ) -> mx.array:
         output = self.mixer(self.norm(x), mask, cache)
         return output + x
@@ -196,7 +215,7 @@ class Mamba2(nn.Module):
         self.norm_f = nn.RMSNorm(args.hidden_size, eps=args.layer_norm_epsilon)
 
     def __call__(
-        self, x: mx.array, cache: Optional[list[MambaCache]] = None
+        self, x: mx.array, cache: Optional[list[ArraysCache]] = None
     ) -> mx.array:
         hidden = self.embeddings(x)
 
@@ -221,7 +240,7 @@ class Model(nn.Module):
             self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
 
     def __call__(
-        self, inputs: mx.array, cache: Optional[list[MambaCache]] = None
+        self, inputs: mx.array, cache: Optional[list[ArraysCache]] = None
     ) -> mx.array:
         hidden = self.backbone(inputs, cache)
 
@@ -231,8 +250,8 @@ class Model(nn.Module):
             logits = self.lm_head(hidden)
         return logits
 
-    def make_cache(self, batch_size: int = 1) -> list[MambaCache]:
-        return [MambaCache() for _ in range(self.args.num_hidden_layers)]
+    def make_cache(self, batch_size: int = 1) -> list[ArraysCache]:
+        return [ArraysCache(size=2) for _ in range(self.args.num_hidden_layers)]
 
     @property
     def layers(self):

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

@@ -1,10 +1,12 @@
 # Copyright © 2025 Apple Inc.
 
 from dataclasses import dataclass
+from functools import lru_cache
 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
@@ -32,6 +34,55 @@ class ModelArgs(BaseModelArgs):
     use_qk_norm: bool = True
 
 
+@lru_cache
+def sharded_rms_norm(group):
+    @mx.compile
+    def _cast_square_sum(x):
+        return x.astype(mx.float32).square().sum(-1, keepdims=True)
+
+    @mx.compile
+    def _normalize(x, norm2, w, eps):
+        norm2 = mx.distributed.all_sum(norm2, group=group)
+        norm = mx.rsqrt(norm2 / (x.shape[-1] * group.size()) + eps)
+        return (x.astype(mx.float32) * norm * w).astype(x.dtype)
+
+    # Split the compile so that x upcasting doesn't break the compile and we
+    # have 2 kernels generated 1 for f(x) = square(upcast(x)) and another
+    # g(x) = downcast(upcast(x) * norm * w)
+    def _inner_sharded_rms_norm(x, w, eps):
+        return _normalize(x, _cast_square_sum(x), w, eps)
+
+    return _inner_sharded_rms_norm
+
+
+class ShardedRMSNorm(nn.Module):
+    def __init__(
+        self, dims: int, eps: float = 1e-5, group: Optional[mx.distributed.Group] = None
+    ):
+        super().__init__()
+        group = group or mx.distributed.init()
+        self.weight = mx.ones((dims // group.size(),))
+        self.group = group
+        self.eps = eps
+
+    def _extra_repr(self):
+        return f"{self.weight.shape[0] * self.group.size()}, eps={self.eps}"
+
+    def __call__(self, x):
+        return sharded_rms_norm(self.group)(x, self["weight"], self.eps)
+
+    @classmethod
+    def from_rms_norm(
+        cls, norm_module, *, group: Optional[mx.distributed.Group] = None
+    ):
+        sn = cls(norm_module.weight.shape[0], norm_module.eps, group=group)
+        sn.weight = mx.contiguous(
+            mx.split(norm_module.weight, group.size(), axis=-1)[group.rank()]
+        )
+
+        return sn
+
+
 class MiniMaxAttention(nn.Module):
     def __init__(self, args: ModelArgs):
         super().__init__()
@@ -118,8 +169,12 @@ class MiniMaxSparseMoeBlock(nn.Module):
             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)
@@ -135,6 +190,10 @@ class MiniMaxSparseMoeBlock(nn.Module):
 
         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
 
 
@@ -218,7 +277,8 @@ class Model(nn.Module):
         """Dequantize FP8 weights and restructure MoE experts."""
 
         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
@@ -266,6 +326,53 @@ class Model(nn.Module):
 
         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 = ShardedRMSNorm.from_rms_norm(
+                    layer.self_attn.q_norm, group=group
+                )
+                layer.self_attn.k_norm = ShardedRMSNorm.from_rms_norm(
+                    layer.self_attn.k_norm, group=group
+                )
+
+            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

mlx_lm/models/ministral3.py

@@ -5,9 +5,12 @@ 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
 
 
@@ -36,13 +39,17 @@ class ModelArgs(BaseModelArgs):
             self.layer_types = ["full_attention"] * self.num_hidden_layers
 
 
-def _get_llama_4_attn_scale(
-    start: int, stop: int, beta: float, max_position_embeddings: int
-):
+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(start, stop) / max_position_embeddings)
+        1 + mx.floor((mx.arange(size) + offset) / max_position_embeddings)
     )
-    return scaling[:, None]
+    if scaling.ndim == 2:
+        return scaling[:, None, :, None]
+    else:
+        return scaling[:, None]
 
 
 class Attention(nn.Module):
@@ -115,7 +122,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):
@@ -146,7 +153,7 @@ class TransformerBlock(nn.Module):
         return out
 
 
-class LanguageModel(nn.Module):
+class LanguageModel(PipelineMixin, nn.Module):
     def __init__(self, args: ModelArgs):
         super().__init__()
         self.args = args
@@ -167,6 +174,18 @@ class LanguageModel(nn.Module):
                 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,
@@ -178,28 +197,47 @@ class LanguageModel(nn.Module):
         else:
             h = self.embed_tokens(inputs)
 
+        pipeline_rank = self.pipeline_rank
+        pipeline_size = self.pipeline_size
+
         if cache is None:
-            cache = [None] * len(self.layers)
+            cache = [None] * len(self.pipeline_layers)
             offset = 0
         else:
             offset = cache[0].offset
 
-        fa_mask = create_attention_mask(h, cache[self.fa_idx])
+        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,
-            offset + inputs.shape[1],
             self.args.rope_parameters["llama_4_scaling_beta"],
             self.args.rope_parameters["original_max_position_embeddings"],
         ).astype(h.dtype)
 
-        for layer, cache in zip(self.layers, cache):
-            mask = swa_mask if layer.use_sliding else fa_mask
-            h = layer(h, attn_scale, mask, cache=cache)
+        # 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)
 
@@ -249,9 +287,40 @@ class Model(nn.Module):
 
         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
+        return self.model.pipeline_layers
 
     def make_cache(self):
         return [

mlx_lm/models/mixtral.py

@@ -1,4 +1,4 @@
-# Copyright © 2023-2024 Apple Inc.
+# Copyright © 2026 Apple Inc.
 
 from dataclasses import dataclass
 from typing import Any, Dict, Optional, Union
@@ -25,6 +25,7 @@ class ModelArgs(BaseModelArgs):
     rope_theta: float = 1e6
     rope_traditional: bool = False
     rope_scaling: Optional[Dict[str, Union[float, str]]] = None
+    tie_word_embeddings: bool = False
 
     def __post_init__(self):
         if self.num_key_value_heads is None:
@@ -162,8 +163,12 @@ class MixtralModel(nn.Module):
         self,
         inputs: mx.array,
         cache=None,
-    ):
-        h = self.embed_tokens(inputs)
+        input_embeddings: Optional[mx.array] = None,
+    ) -> mx.array:
+        if input_embeddings is not None:
+            h = input_embeddings
+        else:
+            h = self.embed_tokens(inputs)
 
         if cache is None:
             cache = [None] * len(self.layers)
@@ -179,20 +184,27 @@ class MixtralModel(nn.Module):
 class Model(nn.Module):
     def __init__(self, args: ModelArgs):
         super().__init__()
+        self.args = args
         self.model_type = args.model_type
         self.model = MixtralModel(args)
-        self.lm_head = nn.Linear(args.hidden_size, args.vocab_size, bias=False)
-        self.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: mx.array,
         cache=None,
-    ):
-        out = self.model(inputs, cache)
-        return self.lm_head(out)
+        input_embeddings: Optional[mx.array] = None,
+    ) -> mx.array:
+        out = self.model(inputs, cache, input_embeddings)
+        if self.args.tie_word_embeddings:
+            return self.model.embed_tokens.as_linear(out)
+        else:
+            return self.lm_head(out)
 
     def sanitize(self, weights):
+        if self.args.tie_word_embeddings:
+            weights.pop("lm_head.weight", None)
         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):

mlx_lm/models/mla.py

@@ -0,0 +1,85 @@
+# Copyright © 2026 Apple Inc.
+
+import math
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+class MultiLinear(nn.Module):
+    def __init__(self, input_dims: int, output_dims: int, num_heads: int) -> None:
+        super().__init__()
+        scale = math.sqrt(1.0 / input_dims)
+        self.weight = mx.random.uniform(
+            low=-scale,
+            high=scale,
+            shape=(num_heads, output_dims, input_dims),
+        )
+
+    def __call__(self, x, transpose=True):
+        if transpose:
+            return x @ self.weight.swapaxes(-1, -2)
+        else:
+            return x @ self.weight
+
+    def to_quantized(
+        self,
+        group_size: int,
+        bits: int,
+        mode: str = "affine",
+    ):
+        num_heads, output_dims, input_dims = self.weight.shape
+        ql = QuantizedMultiLinear(
+            input_dims, output_dims, num_heads, group_size, bits, mode
+        )
+        ql.weight, ql.scales, *biases = mx.quantize(
+            self.weight,
+            group_size,
+            bits,
+            mode=mode,
+        )
+        ql.biases = biases[0] if biases else None
+        return ql
+
+
+class QuantizedMultiLinear(nn.Module):
+    def __init__(
+        self,
+        input_dims: int,
+        output_dims: int,
+        num_heads: int,
+        group_size: int,
+        bits: int,
+        mode: str,
+    ):
+        super().__init__()
+
+        self.group_size = group_size
+        self.bits = bits
+        self.mode = mode
+
+        # Initialize the quantized weight
+        scale = math.sqrt(1 / input_dims)
+        weight = mx.random.uniform(
+            low=-scale,
+            high=scale,
+            shape=(num_heads, output_dims, input_dims),
+        )
+        self.weight, self.scales, *biases = mx.quantize(
+            weight, group_size, bits, mode=mode
+        )
+        self.biases = biases[0] if biases else None
+
+        self.freeze()
+
+    def __call__(self, x, transpose=True):
+        return mx.quantized_matmul(
+            x,
+            self["weight"],
+            scales=self["scales"],
+            biases=self.get("biases"),
+            transpose=transpose,
+            group_size=self.group_size,
+            bits=self.bits,
+            mode=self.mode,
+        )

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,14 +7,16 @@ 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,
     create_ssm_mask,
     scaled_dot_product_attention,
 )
-from .cache import KVCache, MambaCache
+from .cache import ArraysCache, KVCache
 from .ssm import ssm_update
+from .switch_layers import SwitchMLP
 
 
 @dataclass()
@@ -34,27 +36,56 @@ class ModelArgs(BaseModelArgs):
     ssm_state_size: int
     conv_kernel: int
     n_groups: int
-    time_step_limit: Tuple[float, float]
     mlp_bias: bool
     layer_norm_epsilon: float
-    rms_norm_eps: float
     use_bias: bool
     use_conv_bias: bool
-    residual_in_fp32: bool
-    hybrid_override_pattern: List[str]
+    hybrid_override_pattern: Optional[List[str]] = None
+    layers_block_type: Optional[List[str]] = None
     head_dim: Optional[int] = None
+    moe_intermediate_size: Optional[int] = None
+    moe_shared_expert_intermediate_size: Optional[int] = None
+    moe_latent_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
+    time_step_limit: Optional[Tuple[float, float]] = None
+    time_step_min: Optional[float] = None
+    time_step_max: Optional[float] = None
+
+    # Map from layers_block_type names to single-char pattern codes
+    _block_type_to_char = {"mamba": "M", "attention": "*", "moe": "E", "mlp": "-"}
+
+    def __post_init__(self):
+        if self.time_step_limit is None:
+            self.time_step_limit = (0.0, float("inf"))
+
+        # Normalize to hybrid_override_pattern (single-char list)
+        if self.hybrid_override_pattern is None and self.layers_block_type is not None:
+            self.hybrid_override_pattern = [
+                self._block_type_to_char[t] for t in self.layers_block_type
+            ]
+        if self.hybrid_override_pattern is not None:
+            self.num_hidden_layers = len(self.hybrid_override_pattern)
 
 
 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 +121,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[ArraysCache],
+        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 +149,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 +171,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[ArraysCache],
+        mask: Optional[mx.array],
     ) -> mx.array:
         batch_size, seq_len, _ = hidden_states.shape
 
@@ -133,27 +181,34 @@ 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,
         hidden_states: mx.array,
         mask: Optional[mx.array],
-        cache: Optional[MambaCache] = None,
+        cache: Optional[ArraysCache] = None,
     ) -> mx.array:
 
         projected = self.in_proj(hidden_states)
@@ -163,11 +218,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 +227,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 +295,139 @@ 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.moe_latent_size = config.moe_latent_size
+
+        # When latent projection is used, experts operate on the latent dim
+        expert_input_dim = (
+            config.moe_latent_size
+            if config.moe_latent_size is not None
+            else config.hidden_size
+        )
+        self.switch_mlp = SwitchMLP(
+            expert_input_dim,
+            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
+            )
+
+        # Latent projection layers for dimensionality reduction before/after experts
+        if config.moe_latent_size is not None:
+            self.fc1_latent_proj = nn.Linear(
+                config.hidden_size, config.moe_latent_size, bias=config.mlp_bias
+            )
+            self.fc2_latent_proj = nn.Linear(
+                config.moe_latent_size, config.hidden_size, bias=config.mlp_bias
+            )
+
+    def __call__(self, x):
+        residuals = x
+        inds, scores = self.gate(x)
+
+        if self.moe_latent_size is not None:
+            x = self.fc1_latent_proj(x)
+
+        y = self.switch_mlp(x, inds)
+        y = (y * scores[..., None]).sum(axis=-2).astype(y.dtype)
+
+        if self.moe_latent_size is not None:
+            y = self.fc2_latent_proj(y)
+
+        if self.config.n_shared_experts is not None:
+            y = y + self.shared_experts(residuals)
+
+        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 +437,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 +463,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:
@@ -363,13 +530,33 @@ class Model(nn.Module):
         caches = []
         for l in self.layers:
             if l.block_type == "M":
-                caches.append(MambaCache())
+                caches.append(ArraysCache(size=2))
             elif l.block_type == "*":
                 caches.append(KVCache())
         return caches
 
     def sanitize(self, weights):
+        weights = {k: v for (k, v) in weights.items() if not k.startswith("mtp.")}
         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

@@ -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
 from .cache import KVCache, RotatingKVCache
 from .rope_utils import initialize_rope
@@ -131,7 +132,7 @@ class Olmo3MLP(nn.Module):
         self.up_proj = nn.Linear(args.hidden_size, args.intermediate_size, bias=False)
 
     def __call__(self, x: mx.array) -> mx.array:
-        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+        return self.down_proj(swiglu(self.gate_proj(x), self.up_proj(x)))
 
 
 class Olmo3DecoderLayer(nn.Module):

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/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):

mlx_lm/models/plamo2.py

@@ -9,7 +9,8 @@ import mlx.nn as nn
 
 from mlx_lm.models.base import BaseModelArgs, create_attention_mask, create_ssm_mask
 
-from .cache import KVCache, MambaCache
+from .activations import swiglu
+from .cache import ArraysCache, KVCache
 from .ssm import ssm_update
 
 
@@ -54,27 +55,13 @@ class RMSNorm(nn.Module):
         )
 
 
-def causal_conv1d_update(conv_state, x, weight) -> tuple[mx.array, mx.array]:
-    dim = x.shape[-1]
-    state_len = conv_state.shape[-2]
-    x = mx.concatenate([conv_state, x], axis=-2)
-    conv_state = x[:, -state_len:]
-    out = mx.conv1d(
-        x,
-        weight,
-        padding=0,
-        groups=dim,
-    )
-    return nn.silu(out), conv_state
-
-
 class Mamba(nn.Module):
     def __init__(self, config: ModelArgs) -> None:
         super().__init__()
         self.config = config
         self.hidden_size = config.hidden_size
         self.d_state = config.mamba_d_state
-        self.d_conv = config.mamba_d_conv
+        self.conv_kernel_size = config.mamba_d_conv
         self.chunk_size = config.mamba_chunk_size
         self.num_heads = config.mamba_num_heads
         self.hidden_size_per_head = config.hidden_size_per_head
@@ -88,7 +75,7 @@ class Mamba(nn.Module):
             in_channels=self.intermediate_size,
             out_channels=self.intermediate_size,
             bias=False,
-            kernel_size=self.d_conv,
+            kernel_size=self.conv_kernel_size,
             groups=self.intermediate_size,
             padding=0,
         )
@@ -111,20 +98,63 @@ class Mamba(nn.Module):
 
         self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
 
+    def _conv(
+        self,
+        conv_input: mx.array,
+        cache: Optional[ArraysCache],
+        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.intermediate_size,
+                    ),
+                    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,
         x: mx.array,
         B: mx.array,
         C: mx.array,
         dt: mx.array,
-        state: Optional[mx.array] = None,
-        mask: Optional[mx.array] = None,
+        cache: Optional[Any],
+        mask: Optional[mx.array],
     ) -> mx.array:
         batch_size, seq_len, _ = x.shape
 
         x = x.reshape(batch_size, seq_len, self.num_heads, self.hidden_size_per_head)
         B = B.reshape(batch_size, seq_len, 1, self.d_state)
         C = C.reshape(batch_size, seq_len, 1, self.d_state)
+        if cache:
+            state = cache[1]
+            lengths = cache.lengths
+        else:
+            state, lengths = None, None
 
         y, state = ssm_update(
             x,
@@ -136,8 +166,11 @@ class Mamba(nn.Module):
             self.dt_bias,
             state,
             mask=mask,
+            lengths=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,
@@ -147,14 +180,6 @@ class Mamba(nn.Module):
     ):
         bsize, length, _ = hidden_states.shape
 
-        if cache is not None and cache[0] is not None:
-            conv_state = cache[0]
-        else:
-            conv_state = mx.zeros(
-                (bsize, self.d_conv - 1, self.intermediate_size),
-                dtype=hidden_states.dtype,
-            )
-
         zx = self.in_proj(hidden_states)
         zx = zx.reshape(bsize, length, self.num_heads, -1)
         # z: (bsize, length, num_heads, hidden_size_per_head)
@@ -168,9 +193,8 @@ class Mamba(nn.Module):
         )
 
         x = x.reshape(bsize, -1, self.num_heads * self.hidden_size_per_head)
-        if mask is not None:
-            x = mx.where(mask[..., None], x, 0)
-        x, conv_state = causal_conv1d_update(conv_state, x, self.conv1d.weight)
+        x = self._conv(x, cache, mask)
+
         BCdt = self.bcdt_proj(x)
         B, C, dt = mx.split(BCdt, [self.d_state, self.d_state * 2], axis=-1)
 
@@ -181,18 +205,18 @@ class Mamba(nn.Module):
 
         # (bsize, length, num_heads)
         dt = self.dt_proj(dt)
-        out, ssm_state = self._ssm(
+        out = self._ssm(
             x,
             B,
             C,
             dt,
-            cache[1] if cache else None,
+            cache,
             mask,
         )
-        out = out * nn.silu(z.flatten(-2))
-        if cache is not None:
-            cache[0] = conv_state
-            cache[1] = ssm_state
+        if cache:
+            cache.advance(out.shape[1])
+
+        out = swiglu(z.flatten(-2), out)
         return self.out_proj(out)
 
 
@@ -282,7 +306,7 @@ class MLP(nn.Module):
     def __call__(self, x: mx.array) -> mx.array:
         h = self.gate_up_proj(x)
         hs = mx.split(h, 2, axis=-1)
-        return self.down_proj(nn.silu(hs[0]) * hs[1])
+        return self.down_proj(swiglu(hs[0], hs[1]))
 
 
 class PlamoDecoderLayer(nn.Module):
@@ -435,7 +459,7 @@ class Model(nn.Module):
     def make_cache(self):
         # TODO use RotatingKVCache is not full_attn
         # full_attn = self.layer_idx in self.config.full_attention_idx
-        return [MambaCache() if l.is_mamba else KVCache() for l in self.layers]
+        return [ArraysCache(size=2) if l.is_mamba else KVCache() for l in self.layers]
 
     def __call__(self, inputs: mx.array, cache=None) -> mx.array:
         outputs = self.model(

mlx_lm/models/qwen.py

@@ -5,6 +5,7 @@ from dataclasses import dataclass
 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
 
 
@@ -89,7 +90,7 @@ class MLP(nn.Module):
     def __call__(self, x):
         a1 = self.w1(x)
         a2 = self.w2(x)
-        return self.c_proj(a1 * nn.silu(a2))
+        return self.c_proj(swiglu(a2, a1))
 
 
 class TransformerBlock(nn.Module):

mlx_lm/models/qwen2.py

@@ -5,7 +5,9 @@ from typing import Any, Dict, 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 .rope_utils import initialize_rope
 
@@ -90,7 +92,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):
@@ -183,6 +185,37 @@ class Model(nn.Module):
             k: v for k, v in weights.items() if "self_attn.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
+            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

mlx_lm/models/qwen2_moe.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 .switch_layers import SwitchGLU
 
@@ -103,7 +104,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 Qwen2MoeSparseMoeBlock(nn.Module):